From fe50e8d1df075abec8a265d5b164b085de792bb3 Mon Sep 17 00:00:00 2001 From: dxqb Date: Sun, 10 May 2026 14:06:47 +0200 Subject: [PATCH 01/21] feat: split CTK UI into controller/base-view/ctk-view pattern Co-Authored-By: Claude Sonnet 4.6 --- ....py => BaseAdditionalEmbeddingsTabView.py} | 0 .../ui/{CaptionUI.py => BaseCaptionUIView.py} | 0 .../ui/{CloudTab.py => BaseCloudTabView.py} | 0 .../{ConceptTab.py => BaseConceptTabView.py} | 0 ...ceptWindow.py => BaseConceptWindowView.py} | 0 .../{ConfigList.py => BaseConfigListView.py} | 0 ...rtModelUI.py => BaseConvertModelUIView.py} | 0 ...w.py => BaseGenerateCaptionsWindowView.py} | 0 ...ndow.py => BaseGenerateMasksWindowView.py} | 0 modules/ui/{LoraTab.py => BaseLoraTabView.py} | 0 .../ui/{ModelTab.py => BaseModelTabView.py} | 0 ...damWindow.py => BaseMuonAdamWindowView.py} | 0 ...gWindow.py => BaseOffloadingWindowView.py} | 0 ...ow.py => BaseOptimizerParamsWindowView.py} | 0 ...ngWindow.py => BaseProfilingWindowView.py} | 0 ...{SampleFrame.py => BaseSampleFrameView.py} | 0 ...indow.py => BaseSampleParamsWindowView.py} | 0 ...ampleWindow.py => BaseSampleWindowView.py} | 0 ...{SamplingTab.py => BaseSamplingTabView.py} | 0 ...ow.py => BaseSchedulerParamsWindowView.py} | 0 ... => BaseTimestepDistributionWindowView.py} | 0 modules/ui/{TopBar.py => BaseTopBarView.py} | 0 modules/ui/{TrainUI.py => BaseTrainUIView.py} | 0 ...{TrainingTab.py => BaseTrainingTabView.py} | 0 ...{VideoToolUI.py => BaseVideoToolUIView.py} | 0 modules/ui/CaptionUIController.py | 572 +++++++++++ modules/ui/ConceptWindowController.py | 934 ++++++++++++++++++ modules/ui/ConvertModelUIController.py | 170 ++++ modules/ui/CtkAdditionalEmbeddingsTabView.py | 136 +++ modules/ui/CtkCaptionUIView.py | 572 +++++++++++ modules/ui/CtkCloudTabView.py | 221 +++++ modules/ui/CtkConceptTabView.py | 286 ++++++ modules/ui/CtkConceptWindowView.py | 934 ++++++++++++++++++ modules/ui/CtkConfigListView.py | 354 +++++++ modules/ui/CtkConvertModelUIView.py | 170 ++++ modules/ui/CtkGenerateCaptionsWindowView.py | 133 +++ modules/ui/CtkGenerateMasksWindowView.py | 151 +++ modules/ui/CtkLoraTabView.py | 154 +++ modules/ui/CtkModelTabView.py | 688 +++++++++++++ modules/ui/CtkMuonAdamWindowView.py | 107 ++ modules/ui/CtkOffloadingWindowView.py | 75 ++ modules/ui/CtkOptimizerParamsWindowView.py | 288 ++++++ modules/ui/CtkProfilingWindowView.py | 57 ++ modules/ui/CtkSampleFrameView.py | 134 +++ modules/ui/CtkSampleParamsWindowView.py | 39 + modules/ui/CtkSampleWindowView.py | 227 +++++ modules/ui/CtkSamplingTabView.py | 124 +++ modules/ui/CtkSchedulerParamsWindowView.py | 119 +++ .../ui/CtkTimestepDistributionWindowView.py | 186 ++++ modules/ui/CtkTopBarView.py | 260 +++++ modules/ui/CtkTrainUIView.py | 889 +++++++++++++++++ modules/ui/CtkTrainingTabView.py | 856 ++++++++++++++++ modules/ui/CtkVideoToolUIView.py | 877 ++++++++++++++++ .../ui/GenerateCaptionsWindowController.py | 133 +++ modules/ui/GenerateMasksWindowController.py | 151 +++ modules/ui/OptimizerParamsWindowController.py | 288 ++++++ modules/ui/SampleWindowController.py | 227 +++++ .../TimestepDistributionWindowController.py | 186 ++++ modules/ui/TopBarController.py | 260 +++++ modules/ui/TrainUIController.py | 889 +++++++++++++++++ modules/ui/VideoToolUIController.py | 877 ++++++++++++++++ modules/util/ui/ctk_validation.py | 501 ++++++++++ 62 files changed, 13225 insertions(+) rename modules/ui/{AdditionalEmbeddingsTab.py => BaseAdditionalEmbeddingsTabView.py} (100%) rename modules/ui/{CaptionUI.py => BaseCaptionUIView.py} (100%) rename modules/ui/{CloudTab.py => BaseCloudTabView.py} (100%) rename modules/ui/{ConceptTab.py => BaseConceptTabView.py} (100%) rename modules/ui/{ConceptWindow.py => BaseConceptWindowView.py} (100%) rename modules/ui/{ConfigList.py => BaseConfigListView.py} (100%) rename modules/ui/{ConvertModelUI.py => BaseConvertModelUIView.py} (100%) rename modules/ui/{GenerateCaptionsWindow.py => BaseGenerateCaptionsWindowView.py} (100%) rename modules/ui/{GenerateMasksWindow.py => BaseGenerateMasksWindowView.py} (100%) rename modules/ui/{LoraTab.py => BaseLoraTabView.py} (100%) rename modules/ui/{ModelTab.py => BaseModelTabView.py} (100%) rename modules/ui/{MuonAdamWindow.py => BaseMuonAdamWindowView.py} (100%) rename modules/ui/{OffloadingWindow.py => BaseOffloadingWindowView.py} (100%) rename modules/ui/{OptimizerParamsWindow.py => BaseOptimizerParamsWindowView.py} (100%) rename modules/ui/{ProfilingWindow.py => BaseProfilingWindowView.py} (100%) rename modules/ui/{SampleFrame.py => BaseSampleFrameView.py} (100%) rename modules/ui/{SampleParamsWindow.py => BaseSampleParamsWindowView.py} (100%) rename modules/ui/{SampleWindow.py => BaseSampleWindowView.py} (100%) rename modules/ui/{SamplingTab.py => BaseSamplingTabView.py} (100%) rename modules/ui/{SchedulerParamsWindow.py => BaseSchedulerParamsWindowView.py} (100%) rename modules/ui/{TimestepDistributionWindow.py => BaseTimestepDistributionWindowView.py} (100%) rename modules/ui/{TopBar.py => BaseTopBarView.py} (100%) rename modules/ui/{TrainUI.py => BaseTrainUIView.py} (100%) rename modules/ui/{TrainingTab.py => BaseTrainingTabView.py} (100%) rename modules/ui/{VideoToolUI.py => BaseVideoToolUIView.py} (100%) create mode 100644 modules/ui/CaptionUIController.py create mode 100644 modules/ui/ConceptWindowController.py create mode 100644 modules/ui/ConvertModelUIController.py create mode 100644 modules/ui/CtkAdditionalEmbeddingsTabView.py create mode 100644 modules/ui/CtkCaptionUIView.py create mode 100644 modules/ui/CtkCloudTabView.py create mode 100644 modules/ui/CtkConceptTabView.py create mode 100644 modules/ui/CtkConceptWindowView.py create mode 100644 modules/ui/CtkConfigListView.py create mode 100644 modules/ui/CtkConvertModelUIView.py create mode 100644 modules/ui/CtkGenerateCaptionsWindowView.py create mode 100644 modules/ui/CtkGenerateMasksWindowView.py create mode 100644 modules/ui/CtkLoraTabView.py create mode 100644 modules/ui/CtkModelTabView.py create mode 100644 modules/ui/CtkMuonAdamWindowView.py create mode 100644 modules/ui/CtkOffloadingWindowView.py create mode 100644 modules/ui/CtkOptimizerParamsWindowView.py create mode 100644 modules/ui/CtkProfilingWindowView.py create mode 100644 modules/ui/CtkSampleFrameView.py create mode 100644 modules/ui/CtkSampleParamsWindowView.py create mode 100644 modules/ui/CtkSampleWindowView.py create mode 100644 modules/ui/CtkSamplingTabView.py create mode 100644 modules/ui/CtkSchedulerParamsWindowView.py create mode 100644 modules/ui/CtkTimestepDistributionWindowView.py create mode 100644 modules/ui/CtkTopBarView.py create mode 100644 modules/ui/CtkTrainUIView.py create mode 100644 modules/ui/CtkTrainingTabView.py create mode 100644 modules/ui/CtkVideoToolUIView.py create mode 100644 modules/ui/GenerateCaptionsWindowController.py create mode 100644 modules/ui/GenerateMasksWindowController.py create mode 100644 modules/ui/OptimizerParamsWindowController.py create mode 100644 modules/ui/SampleWindowController.py create mode 100644 modules/ui/TimestepDistributionWindowController.py create mode 100644 modules/ui/TopBarController.py create mode 100644 modules/ui/TrainUIController.py create mode 100644 modules/ui/VideoToolUIController.py create mode 100644 modules/util/ui/ctk_validation.py diff --git a/modules/ui/AdditionalEmbeddingsTab.py b/modules/ui/BaseAdditionalEmbeddingsTabView.py similarity index 100% rename from modules/ui/AdditionalEmbeddingsTab.py rename to modules/ui/BaseAdditionalEmbeddingsTabView.py diff --git a/modules/ui/CaptionUI.py b/modules/ui/BaseCaptionUIView.py similarity index 100% rename from modules/ui/CaptionUI.py rename to modules/ui/BaseCaptionUIView.py diff --git a/modules/ui/CloudTab.py b/modules/ui/BaseCloudTabView.py similarity index 100% rename from modules/ui/CloudTab.py rename to modules/ui/BaseCloudTabView.py diff --git a/modules/ui/ConceptTab.py b/modules/ui/BaseConceptTabView.py similarity index 100% rename from modules/ui/ConceptTab.py rename to modules/ui/BaseConceptTabView.py diff --git a/modules/ui/ConceptWindow.py b/modules/ui/BaseConceptWindowView.py similarity index 100% rename from modules/ui/ConceptWindow.py rename to modules/ui/BaseConceptWindowView.py diff --git a/modules/ui/ConfigList.py b/modules/ui/BaseConfigListView.py similarity index 100% rename from modules/ui/ConfigList.py rename to modules/ui/BaseConfigListView.py diff --git a/modules/ui/ConvertModelUI.py b/modules/ui/BaseConvertModelUIView.py similarity index 100% rename from modules/ui/ConvertModelUI.py rename to modules/ui/BaseConvertModelUIView.py diff --git a/modules/ui/GenerateCaptionsWindow.py b/modules/ui/BaseGenerateCaptionsWindowView.py similarity index 100% rename from modules/ui/GenerateCaptionsWindow.py rename to modules/ui/BaseGenerateCaptionsWindowView.py diff --git a/modules/ui/GenerateMasksWindow.py b/modules/ui/BaseGenerateMasksWindowView.py similarity index 100% rename from modules/ui/GenerateMasksWindow.py rename to modules/ui/BaseGenerateMasksWindowView.py diff --git a/modules/ui/LoraTab.py b/modules/ui/BaseLoraTabView.py similarity index 100% rename from modules/ui/LoraTab.py rename to modules/ui/BaseLoraTabView.py diff --git a/modules/ui/ModelTab.py b/modules/ui/BaseModelTabView.py similarity index 100% rename from modules/ui/ModelTab.py rename to modules/ui/BaseModelTabView.py diff --git a/modules/ui/MuonAdamWindow.py b/modules/ui/BaseMuonAdamWindowView.py similarity index 100% rename from modules/ui/MuonAdamWindow.py rename to modules/ui/BaseMuonAdamWindowView.py diff --git a/modules/ui/OffloadingWindow.py b/modules/ui/BaseOffloadingWindowView.py similarity index 100% rename from modules/ui/OffloadingWindow.py rename to modules/ui/BaseOffloadingWindowView.py diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/BaseOptimizerParamsWindowView.py similarity index 100% rename from modules/ui/OptimizerParamsWindow.py rename to modules/ui/BaseOptimizerParamsWindowView.py diff --git a/modules/ui/ProfilingWindow.py b/modules/ui/BaseProfilingWindowView.py similarity index 100% rename from modules/ui/ProfilingWindow.py rename to modules/ui/BaseProfilingWindowView.py diff --git a/modules/ui/SampleFrame.py b/modules/ui/BaseSampleFrameView.py similarity index 100% rename from modules/ui/SampleFrame.py rename to modules/ui/BaseSampleFrameView.py diff --git a/modules/ui/SampleParamsWindow.py b/modules/ui/BaseSampleParamsWindowView.py similarity index 100% rename from modules/ui/SampleParamsWindow.py rename to modules/ui/BaseSampleParamsWindowView.py diff --git a/modules/ui/SampleWindow.py b/modules/ui/BaseSampleWindowView.py similarity index 100% rename from modules/ui/SampleWindow.py rename to modules/ui/BaseSampleWindowView.py diff --git a/modules/ui/SamplingTab.py b/modules/ui/BaseSamplingTabView.py similarity index 100% rename from modules/ui/SamplingTab.py rename to modules/ui/BaseSamplingTabView.py diff --git a/modules/ui/SchedulerParamsWindow.py b/modules/ui/BaseSchedulerParamsWindowView.py similarity index 100% rename from modules/ui/SchedulerParamsWindow.py rename to modules/ui/BaseSchedulerParamsWindowView.py diff --git a/modules/ui/TimestepDistributionWindow.py b/modules/ui/BaseTimestepDistributionWindowView.py similarity index 100% rename from modules/ui/TimestepDistributionWindow.py rename to modules/ui/BaseTimestepDistributionWindowView.py diff --git a/modules/ui/TopBar.py b/modules/ui/BaseTopBarView.py similarity index 100% rename from modules/ui/TopBar.py rename to modules/ui/BaseTopBarView.py diff --git a/modules/ui/TrainUI.py b/modules/ui/BaseTrainUIView.py similarity index 100% rename from modules/ui/TrainUI.py rename to modules/ui/BaseTrainUIView.py diff --git a/modules/ui/TrainingTab.py b/modules/ui/BaseTrainingTabView.py similarity index 100% rename from modules/ui/TrainingTab.py rename to modules/ui/BaseTrainingTabView.py diff --git a/modules/ui/VideoToolUI.py b/modules/ui/BaseVideoToolUIView.py similarity index 100% rename from modules/ui/VideoToolUI.py rename to modules/ui/BaseVideoToolUIView.py diff --git a/modules/ui/CaptionUIController.py b/modules/ui/CaptionUIController.py new file mode 100644 index 000000000..e6cc0551e --- /dev/null +++ b/modules/ui/CaptionUIController.py @@ -0,0 +1,572 @@ +import os +import platform +import subprocess +import traceback +from tkinter import filedialog + +from modules.module.Blip2Model import Blip2Model +from modules.module.BlipModel import BlipModel +from modules.module.ClipSegModel import ClipSegModel +from modules.module.MaskByColor import MaskByColor +from modules.module.RembgHumanModel import RembgHumanModel +from modules.module.RembgModel import RembgModel +from modules.module.WDModel import WDModel +from modules.ui.GenerateCaptionsWindow import GenerateCaptionsWindow +from modules.ui.GenerateMasksWindow import GenerateMasksWindow +from modules.util import path_util +from modules.util.image_util import load_image +from modules.util.torch_util import default_device, torch_gc +from modules.util.ui import components +from modules.util.ui.ui_utils import bind_mousewheel, set_window_icon +from modules.util.ui.UIState import UIState + +import torch + +import customtkinter as ctk +import cv2 +import numpy as np +from customtkinter import ScalingTracker, ThemeManager +from PIL import Image, ImageDraw + + +class CaptionUI(ctk.CTkToplevel): + def __init__( + self, + parent, + initial_dir: str | None, + initial_include_subdirectories: bool, + *args, + **kwargs, + ) -> None: + super().__init__(parent, *args, **kwargs) + self.protocol("WM_DELETE_WINDOW", self._on_close) + + self.dir = initial_dir + self.config_ui_data = {"include_subdirectories": initial_include_subdirectories} + self.config_ui_state = UIState(self, self.config_ui_data) + self.image_size = 850 + self.help_text = """ + Keyboard shortcuts when focusing on the prompt input field: + Up arrow: previous image + Down arrow: next image + Return: save + Ctrl+M: only show the mask + Ctrl+D: draw mask editing mode + Ctrl+F: fill mask editing mode + + When editing masks: + Left click: add mask + Right click: remove mask + Mouse wheel: increase or decrease brush size""" + self.masking_model = None + self.captioning_model = None + self.image_rel_paths = [] + self.current_image_index = -1 + self.file_list = None + self.image_labels = [] + self.pil_image = None + self.image_width = 0 + self.image_height = 0 + self.pil_mask = None + self.mask_draw_x = 0 + self.mask_draw_y = 0 + self.mask_draw_radius = 0.01 + self.display_only_mask = False + self.image = None + self.image_label = None + self.mask_editing_mode = 'draw' + self.enable_mask_editing_var = ctk.BooleanVar() + self.mask_editing_alpha = None + self.prompt_var = None + self.prompt_component = None + + + self.title("OneTrainer") + self.geometry("1280x980") + self.resizable(False, False) + + + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=1) + self.grid_columnconfigure(0, weight=1) + + + self.top_bar(self) + + self.bottom_frame = ctk.CTkFrame(self) + self.bottom_frame.grid(row=1, column=0, sticky="nsew") + self.bottom_frame.grid_rowconfigure(0, weight=1) + self.bottom_frame.grid_columnconfigure(0, weight=0) + self.bottom_frame.grid_columnconfigure(1, weight=1) + + self.file_list_column(self.bottom_frame) + self.content_column(self.bottom_frame) + self.load_directory() + + self.wait_visibility() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + def top_bar(self, master): + top_frame = ctk.CTkFrame(master) + top_frame.grid(row=0, column=0, sticky="nsew") + + components.button(top_frame, 0, 0, "Open", self.open_directory, + tooltip="open a new directory") + components.button(top_frame, 0, 1, "Generate Masks", self.open_mask_window, + tooltip="open a dialog to automatically generate masks") + components.button(top_frame, 0, 2, "Generate Captions", self.open_caption_window, + tooltip="open a dialog to automatically generate captions") + + if platform.system() == "Windows": + components.button(top_frame, 0, 3, "Open in Explorer", self.open_in_explorer, + tooltip="open the current image in Explorer") + + components.switch(top_frame, 0, 4, self.config_ui_state, "include_subdirectories", + text="include subdirectories") + + top_frame.grid_columnconfigure(5, weight=1) + + components.button(top_frame, 0, 6, "Help", self.print_help, + tooltip=self.help_text) + + def file_list_column(self, master): + if self.file_list is not None: + self.image_labels = [] + self.file_list.destroy() + + self.file_list = ctk.CTkScrollableFrame(master, width=300) + self.file_list.grid(row=0, column=0, sticky="nsew") + + for i, filename in enumerate(self.image_rel_paths): + def __create_switch_image(index): + def __switch_image(event): + self.switch_image(index) + + return __switch_image + + label = ctk.CTkLabel(self.file_list, text=filename) + label.bind("", __create_switch_image(i)) + + self.image_labels.append(label) + label.grid(row=i, column=0, padx=5, sticky="nsw") + + def content_column(self, master): + image = Image.new("RGBA", (512, 512), (0, 0, 0, 0)) + + right_frame = ctk.CTkFrame(master, fg_color="transparent") + right_frame.grid(row=0, column=1, sticky="nsew") + + right_frame.grid_columnconfigure(4, weight=1) + right_frame.grid_rowconfigure(1, weight=1) + + components.button(right_frame, 0, 0, "Draw", self.draw_mask_editing_mode, + tooltip="draw a mask using a brush") + components.button(right_frame, 0, 1, "Fill", self.fill_mask_editing_mode, + tooltip="draw a mask using a fill tool") + + # checkbox to enable mask editing + self.enable_mask_editing_var = ctk.BooleanVar() + self.enable_mask_editing_var.set(False) + enable_mask_editing_checkbox = ctk.CTkCheckBox( + right_frame, text="Enable Mask Editing", variable=self.enable_mask_editing_var, width=50) + enable_mask_editing_checkbox.grid(row=0, column=2, padx=25, pady=5, sticky="w") + + # mask alpha textbox + self.mask_editing_alpha = ctk.CTkEntry(master=right_frame, width=40, placeholder_text="1.0") + self.mask_editing_alpha.insert(0, "1.0") + self.mask_editing_alpha.grid(row=0, column=3, sticky="e", padx=5, pady=5) + self.bind_key_events(self.mask_editing_alpha) + + mask_editing_alpha_label = ctk.CTkLabel(right_frame, text="Brush Alpha", width=75) + mask_editing_alpha_label.grid(row=0, column=4, padx=0, pady=5, sticky="w") + + # image + self.image = ctk.CTkImage( + light_image=image, + size=(self.image_size, self.image_size) + ) + self.image_label = ctk.CTkLabel( + master=right_frame, text="", image=self.image, height=self.image_size, width=self.image_size + ) + self.image_label.grid(row=1, column=0, columnspan=5, sticky="nsew") + + self.image_label.bind("", self.edit_mask) + self.image_label.bind("", self.edit_mask) + self.image_label.bind("", self.edit_mask) + bind_mousewheel(self.image_label, {self.image_label.children["!label"]}, self.draw_mask_radius) + + # prompt + self.prompt_var = ctk.StringVar() + self.prompt_component = ctk.CTkEntry(right_frame, textvariable=self.prompt_var) + self.prompt_component.grid(row=2, column=0, columnspan=5, pady=5, sticky="new") + self.bind_key_events(self.prompt_component) + self.prompt_component.focus_set() + + def bind_key_events(self, component): + component.bind("", self.next_image) + component.bind("", self.previous_image) + component.bind("", self.save) + component.bind("", self.toggle_mask) + component.bind("", self.draw_mask_editing_mode) + component.bind("", self.fill_mask_editing_mode) + + def load_directory(self, include_subdirectories: bool = False): + self.scan_directory(include_subdirectories) + self.file_list_column(self.bottom_frame) + + if len(self.image_rel_paths) > 0: + self.switch_image(0) + else: + self.switch_image(-1) + + self.prompt_component.focus_set() + + def scan_directory(self, include_subdirectories: bool = False): + def __is_supported_image_extension(filename): + name, ext = os.path.splitext(filename) + return path_util.is_supported_image_extension(ext) and not name.endswith("-masklabel") and not name.endswith("-condlabel") + + self.image_rel_paths = [] + + if not self.dir or not os.path.isdir(self.dir): + return + + if include_subdirectories: + for root, _, files in os.walk(self.dir): + for filename in files: + if __is_supported_image_extension(filename): + self.image_rel_paths.append( + os.path.relpath(os.path.join(root, filename), self.dir) + ) + else: + for _, filename in enumerate(os.listdir(self.dir)): + if __is_supported_image_extension(filename): + self.image_rel_paths.append( + os.path.relpath(os.path.join(self.dir, filename), self.dir) + ) + + def load_image(self): + image_name = "resources/icons/icon.png" + + if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): + image_name = self.image_rel_paths[self.current_image_index] + image_name = os.path.join(self.dir, image_name) + + try: + return load_image(image_name, convert_mode="RGB") + except Exception: + print(f'Could not open image {image_name}') + + def load_mask(self): + if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): + image_name = self.image_rel_paths[self.current_image_index] + mask_name = os.path.splitext(image_name)[0] + "-masklabel.png" + mask_name = os.path.join(self.dir, mask_name) + + try: + return load_image(mask_name, convert_mode='RGB') + except Exception: + return None + else: + return None + + def load_prompt(self): + if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): + image_name = self.image_rel_paths[self.current_image_index] + prompt_name = os.path.splitext(image_name)[0] + ".txt" + prompt_name = os.path.join(self.dir, prompt_name) + + try: + with open(prompt_name, "r", encoding='utf-8') as f: + return f.readlines()[0].strip() + except Exception: + return "" + else: + return "" + + def previous_image(self, event): + if len(self.image_rel_paths) > 0 and (self.current_image_index - 1) >= 0: + self.switch_image(self.current_image_index - 1) + + def next_image(self, event): + if len(self.image_rel_paths) > 0 and (self.current_image_index + 1) < len(self.image_rel_paths): + self.switch_image(self.current_image_index + 1) + + def switch_image(self, index): + if len(self.image_labels) > 0 and self.current_image_index < len(self.image_labels): + self.image_labels[self.current_image_index].configure( + text_color=ThemeManager.theme["CTkLabel"]["text_color"]) + + self.current_image_index = index + if index >= 0: + self.image_labels[index].configure(text_color="#FF0000") + + self.pil_image = self.load_image() + self.pil_mask = self.load_mask() + prompt = self.load_prompt() + + self.image_width = self.pil_image.width + self.image_height = self.pil_image.height + scale = self.image_size / max(self.pil_image.height, self.pil_image.width) + height = int(self.pil_image.height * scale) + width = int(self.pil_image.width * scale) + + self.pil_image = self.pil_image.resize((width, height), Image.Resampling.LANCZOS) + + self.refresh_image() + self.prompt_var.set(prompt) + else: + image = Image.new("RGB", (512, 512), (0, 0, 0)) + self.image.configure(light_image=image) + + def refresh_image(self): + if self.pil_mask: + resized_pil_mask = self.pil_mask.resize( + (self.pil_image.width, self.pil_image.height), + Image.Resampling.NEAREST + ) + + if self.display_only_mask: + self.image.configure(light_image=resized_pil_mask, size=resized_pil_mask.size) + else: + np_image = np.array(self.pil_image).astype(np.float32) / 255.0 + np_mask = np.array(resized_pil_mask).astype(np.float32) / 255.0 + + # normalize mask between 0.3 - 1.0 so we can see image underneath and gauge strength of the alpha + norm_min = 0.3 + np_mask_min = np_mask.min() + if np_mask_min == 0: + # optimize for common case + np_mask = np_mask * (1.0 - norm_min) + norm_min + elif np_mask_min < 1: + # note: min of 1 means we get divide by 0 + np_mask = (np_mask - np_mask_min) / (1.0 - np_mask_min) * (1.0 - norm_min) + norm_min + + np_masked_image = (np_image * np_mask * 255.0).astype(np.uint8) + masked_image = Image.fromarray(np_masked_image, mode='RGB') + + self.image.configure(light_image=masked_image, size=masked_image.size) + else: + self.image.configure(light_image=self.pil_image, size=self.pil_image.size) + + def draw_mask_radius(self, delta, raw_event): + # Wheel up = Increase radius. Wheel down = Decrease radius. + multiplier = 1.0 + (delta * 0.05) + self.mask_draw_radius = max(0.0025, self.mask_draw_radius * multiplier) + + def edit_mask(self, event): + if not self.enable_mask_editing_var.get(): + return + + if event.widget != self.image_label.children["!label"]: + return + + if len(self.image_rel_paths) == 0 or self.current_image_index >= len(self.image_rel_paths): + return + + display_scaling = ScalingTracker.get_window_scaling(self) + + event_x = event.x / display_scaling + event_y = event.y / display_scaling + + start_x = int(event_x / self.pil_image.width * self.image_width) + start_y = int(event_y / self.pil_image.height * self.image_height) + end_x = int(self.mask_draw_x / self.pil_image.width * self.image_width) + end_y = int(self.mask_draw_y / self.pil_image.height * self.image_height) + + self.mask_draw_x = event_x + self.mask_draw_y = event_y + + is_right = False + is_left = False + if event.state & 0x0100 or event.num == 1: # left mouse button + is_left = True + elif event.state & 0x0400 or event.num == 3: # right mouse button + is_right = True + + if self.mask_editing_mode == 'draw': + self.draw_mask(start_x, start_y, end_x, end_y, is_left, is_right) + if self.mask_editing_mode == 'fill': + self.fill_mask(start_x, start_y, end_x, end_y, is_left, is_right) + + def draw_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): + color = None + + adding_to_mask = True + if is_left: + try: + alpha = float(self.mask_editing_alpha.get()) + except Exception: + alpha = 1.0 + rgb_value = int(max(0.0, min(alpha, 1.0)) * 255) # max/min stuff to clamp to 0 - 255 range + color = (rgb_value, rgb_value, rgb_value) + + elif is_right: + color = (0, 0, 0) + adding_to_mask = False + + if color is not None: + if self.pil_mask is None: + if adding_to_mask: + self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(0, 0, 0)) + else: + self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(255, 255, 255)) + + radius = int(self.mask_draw_radius * max(self.pil_mask.width, self.pil_mask.height)) + + draw = ImageDraw.Draw(self.pil_mask) + draw.line((start_x, start_y, end_x, end_y), fill=color, + width=radius + radius + 1) + draw.ellipse((start_x - radius, start_y - radius, + start_x + radius, start_y + radius), fill=color, outline=None) + draw.ellipse((end_x - radius, end_y - radius, end_x + radius, + end_y + radius), fill=color, outline=None) + + self.refresh_image() + + def fill_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): + color = None + + adding_to_mask = True + if is_left: + try: + alpha = float(self.mask_editing_alpha.get()) + except Exception: + alpha = 1.0 + rgb_value = int(max(0.0, min(alpha, 1.0)) * 255) # max/min stuff to clamp to 0 - 255 range + color = (rgb_value, rgb_value, rgb_value) + + elif is_right: + color = (0, 0, 0) + adding_to_mask = False + + if color is not None: + if self.pil_mask is None: + if adding_to_mask: + self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(0, 0, 0)) + else: + self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(255, 255, 255)) + + np_mask = np.array(self.pil_mask).astype(np.uint8) + cv2.floodFill(np_mask, None, (start_x, start_y), color) + self.pil_mask = Image.fromarray(np_mask, 'RGB') + + self.refresh_image() + + def save(self, event): + if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): + image_name = self.image_rel_paths[self.current_image_index] + + prompt_name = os.path.splitext(image_name)[0] + ".txt" + prompt_name = os.path.join(self.dir, prompt_name) + + mask_name = os.path.splitext(image_name)[0] + "-masklabel.png" + mask_name = os.path.join(self.dir, mask_name) + + try: + with open(prompt_name, "w", encoding='utf-8') as f: + f.write(self.prompt_var.get()) + except Exception: + return + + if self.pil_mask: + self.pil_mask.save(mask_name) + + def draw_mask_editing_mode(self, *args): + self.mask_editing_mode = 'draw' + + if args: + # disable default event + return "break" + return None + + def fill_mask_editing_mode(self, *args): + self.mask_editing_mode = 'fill' + + def toggle_mask(self, *args): + self.display_only_mask = not self.display_only_mask + self.refresh_image() + + def open_directory(self): + new_dir = filedialog.askdirectory() + + if new_dir: + self.dir = new_dir + self.load_directory(include_subdirectories=self.config_ui_data["include_subdirectories"]) + + def open_mask_window(self): + dialog = GenerateMasksWindow(self, self.dir, self.config_ui_data["include_subdirectories"]) + self.wait_window(dialog) + self.switch_image(self.current_image_index) + + def open_caption_window(self): + dialog = GenerateCaptionsWindow(self, self.dir, self.config_ui_data["include_subdirectories"]) + self.wait_window(dialog) + self.switch_image(self.current_image_index) + + def open_in_explorer(self): + try: + image_name = self.image_rel_paths[self.current_image_index] + image_name = os.path.realpath(os.path.join(self.dir, image_name)) + subprocess.Popen(f"explorer /select,{image_name}") + except Exception: + traceback.print_exc() + + def load_masking_model(self, model): + model_type = type(self.masking_model).__name__ if self.masking_model else None + + if model == "ClipSeg" and model_type != "ClipSegModel": + self._release_models() + print("loading ClipSeg model, this may take a while") + self.masking_model = ClipSegModel(default_device, torch.float32) + elif model == "Rembg" and model_type != "RembgModel": + self._release_models() + print("loading Rembg model, this may take a while") + self.masking_model = RembgModel(default_device, torch.float32) + elif model == "Rembg-Human" and model_type != "RembgHumanModel": + self._release_models() + print("loading Rembg-Human model, this may take a while") + self.masking_model = RembgHumanModel(default_device, torch.float32) + elif model == "Hex Color" and model_type != "MaskByColor": + self._release_models() + self.masking_model = MaskByColor(default_device, torch.float32) + + def load_captioning_model(self, model): + model_type = type(self.captioning_model).__name__ if self.captioning_model else None + + if model == "Blip" and model_type != "BlipModel": + self._release_models() + print("loading Blip model, this may take a while") + self.captioning_model = BlipModel(default_device, torch.float16) + elif model == "Blip2" and model_type != "Blip2Model": + self._release_models() + print("loading Blip2 model, this may take a while") + self.captioning_model = Blip2Model(default_device, torch.float16) + elif model == "WD14 VIT v2" and model_type != "WDModel": + self._release_models() + print("loading WD14_VIT_v2 model, this may take a while") + self.captioning_model = WDModel(default_device, torch.float16) + + def print_help(self): + print(self.help_text) + + def _release_models(self): + """Release all models from VRAM""" + freed = False + if self.captioning_model is not None: + self.captioning_model = None + freed = True + if self.masking_model is not None: + self.masking_model = None + freed = True + if freed: + torch_gc() + + def _on_close(self): + self._release_models() + self.destroy() + + def destroy(self): + self._release_models() + super().destroy() diff --git a/modules/ui/ConceptWindowController.py b/modules/ui/ConceptWindowController.py new file mode 100644 index 000000000..f58879d5f --- /dev/null +++ b/modules/ui/ConceptWindowController.py @@ -0,0 +1,934 @@ +import fractions +import math +import os +import pathlib +import platform +import random +import threading +import time +import traceback + +from modules.util import concept_stats, path_util +from modules.util.config.ConceptConfig import ConceptConfig +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.BalancingStrategy import BalancingStrategy +from modules.util.enum.ConceptType import ConceptType +from modules.util.image_util import load_image +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +from mgds.LoadingPipeline import LoadingPipeline +from mgds.OutputPipelineModule import OutputPipelineModule +from mgds.PipelineModule import PipelineModule +from mgds.pipelineModules.CapitalizeTags import CapitalizeTags +from mgds.pipelineModules.DropTags import DropTags +from mgds.pipelineModules.RandomBrightness import RandomBrightness +from mgds.pipelineModules.RandomCircularMaskShrink import ( + RandomCircularMaskShrink, +) +from mgds.pipelineModules.RandomContrast import RandomContrast +from mgds.pipelineModules.RandomFlip import RandomFlip +from mgds.pipelineModules.RandomHue import RandomHue +from mgds.pipelineModules.RandomMaskRotateCrop import RandomMaskRotateCrop +from mgds.pipelineModules.RandomRotate import RandomRotate +from mgds.pipelineModules.RandomSaturation import RandomSaturation +from mgds.pipelineModules.ShuffleTags import ShuffleTags +from mgds.pipelineModuleTypes.RandomAccessPipelineModule import ( + RandomAccessPipelineModule, +) + +import torch +from torchvision.transforms import functional + +import customtkinter as ctk +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 + + +class InputPipelineModule( + PipelineModule, + RandomAccessPipelineModule, +): + def __init__(self, data: dict): + super().__init__() + self.data = data + + def length(self) -> int: + return 1 + + def get_inputs(self) -> list[str]: + return [] + + def get_outputs(self) -> list[str]: + return list(self.data.keys()) + + def get_item(self, variation: int, index: int, requested_name: str = None) -> dict: + return self.data + + +class ConceptWindow(ctk.CTkToplevel): + def __init__( + self, + parent, + train_config: TrainConfig, + concept: ConceptConfig, + ui_state: UIState, + image_ui_state: UIState, + text_ui_state: UIState, + *args, **kwargs, + ): + super().__init__(parent, *args, **kwargs) + + self.train_config = train_config + + self.concept = concept + self.ui_state = ui_state + self.image_ui_state = image_ui_state + 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.title("Concept") + self.geometry("800x700") + self.resizable(True, True) + + self.grid_rowconfigure(0, weight=1) + self.grid_columnconfigure(0, weight=1) + + tabview = ctk.CTkTabview(self) + tabview.grid(row=0, column=0, sticky="nsew") + + self.general_tab = self.__general_tab(tabview.add("general"), concept) + self.image_augmentation_tab = self.__image_augmentation_tab(tabview.add("image augmentation")) + self.text_augmentation_tab = self.__text_augmentation_tab(tabview.add("text augmentation")) + self.concept_stats_tab = self.__concept_stats_tab(tabview.add("statistics")) + + #automatic concept scan + self.scan_thread = threading.Thread(target=self.__auto_update_concept_stats, daemon=True) + self.scan_thread.start() + + components.button(self, 1, 0, "ok", self.__ok) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def __general_tab(self, master, concept: ConceptConfig): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, weight=1) + + # name + components.label(frame, 0, 0, "Name", + tooltip="Name of the concept") + components.entry(frame, 0, 1, self.ui_state, "name") + + # enabled + components.label(frame, 1, 0, "Enabled", + tooltip="Enable or disable this concept") + components.switch(frame, 1, 1, self.ui_state, "enabled") + + # concept type + components.label(frame, 2, 0, "Concept Type", + tooltip="STANDARD: Standard finetuning with the sample as training target\n" + "VALIDATION: Use concept for validation instead of training\n" + "PRIOR_PREDICTION: Use the sample to make a prediction using the model as it was before training. This prediction is then used as the training target " + "for the model in training. This can be used as regularisation and to preserve prior model knowledge while finetuning the model on other concepts. " + "Only implemented for LoRA.", + wide_tooltip=True) + components.options(frame, 2, 1, [str(x) for x in list(ConceptType)], self.ui_state, "type") + + # path + components.label(frame, 3, 0, "Path", + tooltip="Path where the training data is located") + components.path_entry(frame, 3, 1, self.ui_state, "path", mode="dir") + 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.path_entry(frame, 4, 2, self.text_ui_state, "prompt_path", mode="file") + + def set_prompt_path_entry_enabled(option: str): + if option == 'concept': + for child in prompt_path_entry.children.values(): + child.configure(state="normal") + else: + for child in prompt_path_entry.children.values(): + child.configure(state="disabled") + + components.options_kv(frame, 4, 1, [ + ("From text file per sample", 'sample'), + ("From single text file", 'concept'), + ("From image file name", 'filename'), + ], self.text_ui_state, "prompt_source", command=set_prompt_path_entry_enabled) + set_prompt_path_entry_enabled(concept.text.prompt_source) + + # include subdirectories + components.label(frame, 5, 0, "Include Subdirectories", + tooltip="Includes images from subdirectories into the dataset") + components.switch(frame, 5, 1, self.ui_state, "include_subdirectories") + + # image variations + components.label(frame, 6, 0, "Image Variations", + tooltip="The number of different image versions to cache if latent caching is enabled.") + components.entry(frame, 6, 1, self.ui_state, "image_variations") + + # text variations + components.label(frame, 7, 0, "Text Variations", + tooltip="The number of different text versions to cache if latent caching is enabled.") + components.entry(frame, 7, 1, self.ui_state, "text_variations") + + # balancing + components.label(frame, 8, 0, "Balancing", + tooltip="The number of samples used during training. Use repeats to multiply the concept, or samples to specify an exact number of samples used in each epoch.") + components.entry(frame, 8, 1, self.ui_state, "balancing") + components.options(frame, 8, 2, [str(x) for x in list(BalancingStrategy)], self.ui_state, "balancing_strategy") + + # loss weight + components.label(frame, 9, 0, "Loss Weight", + tooltip="The loss multiplyer for this concept.") + components.entry(frame, 9, 1, self.ui_state, "loss_weight") + + frame.pack(fill="both", expand=1) + return frame + + def __image_augmentation_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=0) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # header + components.label(frame, 0, 1, "Random", + tooltip="Enable this augmentation with random values") + components.label(frame, 0, 2, "Fixed", + tooltip="Enable this augmentation with fixed values") + + # crop jitter + components.label(frame, 1, 0, "Crop Jitter", + tooltip="Enables random cropping of samples") + components.switch(frame, 1, 1, self.image_ui_state, "enable_crop_jitter") + + # random flip + components.label(frame, 2, 0, "Random Flip", + tooltip="Randomly flip the sample during training") + components.switch(frame, 2, 1, self.image_ui_state, "enable_random_flip") + components.switch(frame, 2, 2, self.image_ui_state, "enable_fixed_flip") + + # random rotation + components.label(frame, 3, 0, "Random Rotation", + tooltip="Randomly rotates the sample during training") + components.switch(frame, 3, 1, self.image_ui_state, "enable_random_rotate") + components.switch(frame, 3, 2, self.image_ui_state, "enable_fixed_rotate") + components.entry(frame, 3, 3, self.image_ui_state, "random_rotate_max_angle") + + # random brightness + components.label(frame, 4, 0, "Random Brightness", + tooltip="Randomly adjusts the brightness of the sample during training") + components.switch(frame, 4, 1, self.image_ui_state, "enable_random_brightness") + components.switch(frame, 4, 2, self.image_ui_state, "enable_fixed_brightness") + components.entry(frame, 4, 3, self.image_ui_state, "random_brightness_max_strength") + + # random contrast + components.label(frame, 5, 0, "Random Contrast", + tooltip="Randomly adjusts the contrast of the sample during training") + components.switch(frame, 5, 1, self.image_ui_state, "enable_random_contrast") + components.switch(frame, 5, 2, self.image_ui_state, "enable_fixed_contrast") + components.entry(frame, 5, 3, self.image_ui_state, "random_contrast_max_strength") + + # random saturation + components.label(frame, 6, 0, "Random Saturation", + tooltip="Randomly adjusts the saturation of the sample during training") + components.switch(frame, 6, 1, self.image_ui_state, "enable_random_saturation") + components.switch(frame, 6, 2, self.image_ui_state, "enable_fixed_saturation") + components.entry(frame, 6, 3, self.image_ui_state, "random_saturation_max_strength") + + # random hue + components.label(frame, 7, 0, "Random Hue", + tooltip="Randomly adjusts the hue of the sample during training") + components.switch(frame, 7, 1, self.image_ui_state, "enable_random_hue") + components.switch(frame, 7, 2, self.image_ui_state, "enable_fixed_hue") + components.entry(frame, 7, 3, self.image_ui_state, "random_hue_max_strength") + + # random circular mask shrink + components.label(frame, 8, 0, "Circular Mask Generation", + tooltip="Automatically create circular masks for masked training") + components.switch(frame, 8, 1, self.image_ui_state, "enable_random_circular_mask_shrink") + + # random rotate and crop + components.label(frame, 9, 0, "Random Rotate and Crop", + tooltip="Randomly rotate the training samples and crop to the masked region") + components.switch(frame, 9, 1, self.image_ui_state, "enable_random_mask_rotate_crop") + + # circular mask generation + components.label(frame, 10, 0, "Resolution Override", + tooltip="Override the resolution for this concept. Optionally specify multiple resolutions separated by a comma, or a single exact resolution in the format x") + components.switch(frame, 10, 2, self.image_ui_state, "enable_resolution_override") + components.entry(frame, 10, 3, self.image_ui_state, "resolution_override") + + # image + image_preview, filename_preview, caption_preview = self.__get_preview_image() + self.image = ctk.CTkImage( + light_image=image_preview, + size=image_preview.size, + ) + image_label = ctk.CTkLabel(master=frame, text="", image=self.image, height=300, width=300) + image_label.grid(row=0, column=4, rowspan=6) + + # refresh preview + update_button_frame = ctk.CTkFrame(master=frame, corner_radius=0, fg_color="transparent") + update_button_frame.grid(row=6, column=4, rowspan=6, sticky="nsew") + update_button_frame.grid_columnconfigure(1, weight=1) + + prev_preview_button = components.button(update_button_frame, 0, 0, "<", command=self.__prev_image_preview) + components.button(update_button_frame, 0, 1, "Update Preview", command=self.__update_image_preview) + next_preview_button = components.button(update_button_frame, 0, 2, ">", command=self.__next_image_preview) + preview_augmentations_switch = ctk.CTkSwitch(update_button_frame, text="Show Augmentations", variable=self.preview_augmentations, command=self.__update_image_preview) + preview_augmentations_switch.grid(row=1, column=0, columnspan=3, padx=5, pady=5) + + prev_preview_button.configure(width=40) + next_preview_button.configure(width=40) + + #caption and filename preview + self.filename_preview = ctk.CTkLabel(master=update_button_frame, text=filename_preview, width=300, anchor="nw", justify="left", padx=10, wraplength=280) + self.filename_preview.grid(row=2, column=0, columnspan=3) + self.caption_preview = ctk.CTkTextbox(master=update_button_frame, width = 300, height = 150, wrap="word", border_width=2) + self.caption_preview.insert(index="1.0", text=caption_preview) + self.caption_preview.configure(state="disabled") + self.caption_preview.grid(row=3, column=0, columnspan=3, rowspan=3) + + frame.pack(fill="both", expand=1) + return frame + + def __text_augmentation_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=0) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # tag shuffling + components.label(frame, 0, 0, "Tag Shuffling", + tooltip="Enables tag shuffling") + components.switch(frame, 0, 1, self.text_ui_state, "enable_tag_shuffling") + + # keep tag count + components.label(frame, 1, 0, "Tag Delimiter", + tooltip="The delimiter between tags") + components.entry(frame, 1, 1, self.text_ui_state, "tag_delimiter") + + # keep tag count + components.label(frame, 2, 0, "Keep Tag Count", + tooltip="The number of tags at the start of the caption that are not shuffled or dropped") + components.entry(frame, 2, 1, self.text_ui_state, "keep_tags_count") + + # tag dropout + components.label(frame, 3, 0, "Tag Dropout", + tooltip="Enables random dropout for tags in the captions.") + components.switch(frame, 3, 1, self.text_ui_state, "tag_dropout_enable") + components.label(frame, 4, 0, "Dropout Mode", + tooltip="Method used to drop captions. 'Full' will drop the entire caption past the 'kept' tags with a certain probability, 'Random' will drop individual tags with the set probability, and 'Random Weighted' will linearly increase the probability of dropping tags, more likely to preseve tags near the front with full probability to drop at the end.") + components.options_kv(frame, 4, 1, [ + ("Full", 'FULL'), + ("Random", 'RANDOM'), + ("Random Weighted", 'RANDOM WEIGHTED'), + ], self.text_ui_state, "tag_dropout_mode", None) + components.label(frame, 4, 2, "Probability", + tooltip="Probability to drop tags, from 0 to 1.") + components.entry(frame, 4, 3, self.text_ui_state, "tag_dropout_probability") + + components.label(frame, 5, 0, "Special Dropout Tags", + tooltip="List of tags which will be whitelisted/blacklisted by dropout. 'Whitelist' tags will never be dropped but all others may be, 'Blacklist' tags may be dropped but all others will never be, 'None' may drop any tags. Can specify either a delimiter-separated list in the field, or a file path to a .txt or .csv file with entries separated by newlines.") + components.options_kv(frame, 5, 1, [ + ("None", 'NONE'), + ("Blacklist", 'BLACKLIST'), + ("Whitelist", 'WHITELIST'), + ], self.text_ui_state, "tag_dropout_special_tags_mode", None) + components.entry(frame, 5, 2, self.text_ui_state, "tag_dropout_special_tags") + components.label(frame, 6, 0, "Special Tags Regex", + tooltip="Interpret special tags with regex, such as 'photo.*' to match 'photo, photograph, photon' but not 'telephoto'. Includes exception for '/(' and '/)' syntax found in many booru/e6 tags.") + components.switch(frame, 6, 1, self.text_ui_state, "tag_dropout_special_tags_regex") + + #capitalization randomization + components.label(frame, 7, 0, "Randomize Capitalization", + tooltip="Enables randomization of capitalization for tags in the caption.") + components.switch(frame, 7, 1, self.text_ui_state, "caps_randomize_enable") + components.label(frame, 7, 2, "Force Lowercase", + tooltip="If enabled, converts the caption to lowercase before any further processing.") + components.switch(frame, 7, 3, self.text_ui_state, "caps_randomize_lowercase") + + components.label(frame, 8, 0, "Captialization Mode", + tooltip="Comma-separated list of types of capitalization randomization to perform. 'capslock' for ALL CAPS, 'title' for First Letter Of Every Word, 'first' for First word only, 'random' for rAndOMiZeD lEtTERs.") + components.entry(frame, 8, 1, self.text_ui_state, "caps_randomize_mode") + components.label(frame, 8, 2, "Probability", + tooltip="Probability to randomize capitialization of each tag, from 0 to 1.") + components.entry(frame, 8, 3, self.text_ui_state, "caps_randomize_probability") + + frame.pack(fill="both", expand=1) + return frame + + def __concept_stats_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0, minsize=150) + frame.grid_columnconfigure(1, weight=0, minsize=150) + frame.grid_columnconfigure(2, weight=0, minsize=150) + frame.grid_columnconfigure(3, weight=0, minsize=150) + + self.cancel_scan_flag = threading.Event() + + #file size + self.file_size_label = components.label(frame, 1, 0, "Total Size", pad=0, + tooltip="Total size of all image, mask, and caption files in MB") + self.file_size_label.configure(font=ctk.CTkFont(underline=True)) + self.file_size_preview = components.label(frame, 2, 0, pad=0, text="-") + + #subdirectory count + self.dir_count_label = components.label(frame, 1, 1, "Directories", pad=0, + tooltip="Total number of directories including and under (if 'include subdirectories' is enabled) the main concept directory") + self.dir_count_label.configure(font=ctk.CTkFont(underline=True)) + self.dir_count_preview = components.label(frame, 2, 1, pad=0, text="-") + + #basic img/vid stats - count of each type in the concept + #the \n at the start of the label gives it better vertical spacing with other rows + self.image_count_label = components.label(frame, 3, 0, "\nTotal Images", pad=0, + tooltip="Total number of image files, any of the extensions " + str(path_util.SUPPORTED_IMAGE_EXTENSIONS) + ", excluding '-masklabel.png and -condlabel.png'") + self.image_count_label.configure(font=ctk.CTkFont(underline=True)) + self.image_count_preview = components.label(frame, 4, 0, pad=0, text="-") + self.video_count_label = components.label(frame, 3, 1, "\nTotal Videos", pad=0, + tooltip="Total number of video files, any of the extensions " + str(path_util.SUPPORTED_VIDEO_EXTENSIONS)) + self.video_count_label.configure(font=ctk.CTkFont(underline=True)) + self.video_count_preview = components.label(frame, 4, 1, pad=0, text="-") + self.mask_count_label = components.label(frame, 3, 2, "\nTotal Masks", pad=0, + tooltip="Total number of mask files, any file ending in '-masklabel.png'") + self.mask_count_label.configure(font=ctk.CTkFont(underline=True)) + self.mask_count_preview = components.label(frame, 4, 2, pad=0, text="-") + self.caption_count_label = components.label(frame, 3, 3, "\nTotal Captions", pad=0, + tooltip="Total number of caption files, any .txt file. With advanced scan, includes the total number of captions on separate lines across all files in parentheses.") + self.caption_count_label.configure(font=ctk.CTkFont(underline=True)) + self.caption_count_preview = components.label(frame, 4, 3, pad=0, text="-") + + #advanced img/vid stats - how many img/vid files have a mask or caption of the same name + self.image_count_mask_label = components.label(frame, 5, 0, "\nImages with Masks", pad=0, + tooltip="Total number of image files with an associated mask") + self.image_count_mask_label.configure(font=ctk.CTkFont(underline=True)) + self.image_count_mask_preview = components.label(frame, 6, 0, pad=0, text="-") + self.mask_count_label_unpaired = components.label(frame, 5, 1, "\nUnpaired Masks", pad=0, + tooltip="Total number of mask files which lack a corresponding image file - if >0, check your data set!") + self.mask_count_label_unpaired.configure(font=ctk.CTkFont(underline=True)) + self.mask_count_preview_unpaired = components.label(frame, 6, 1, pad=0, text="-") + #currently no masks for videos? + + self.image_count_caption_label = components.label(frame, 7, 0, "\nImages with Captions", pad=0, + tooltip="Total number of image files with an associated caption") + self.image_count_caption_label.configure(font=ctk.CTkFont(underline=True)) + self.image_count_caption_preview = components.label(frame, 8, 0, pad=0, text="-") + self.video_count_caption_label = components.label(frame, 7, 1, "\nVideos with Captions", pad=0, + tooltip="Total number of video files with an associated caption") + self.video_count_caption_label.configure(font=ctk.CTkFont(underline=True)) + self.video_count_caption_preview = components.label(frame, 8, 1, pad=0, text="-") + self.caption_count_label_unpaired = components.label(frame, 7, 2, "\nUnpaired Captions", pad=0, + tooltip="Total number of caption files which lack a corresponding image file - if >0, check your data set! If using 'from file name' or 'from single text file' then this can be ignored.") + self.caption_count_label_unpaired.configure(font=ctk.CTkFont(underline=True)) + self.caption_count_preview_unpaired = components.label(frame, 8, 2, pad=0, text="-") + + #resolution info + self.pixel_max_label = components.label(frame, 9, 0, "\nMax Pixels", pad=0, + tooltip="Largest image in the concept by number of pixels (width * height)") + self.pixel_max_label.configure(font=ctk.CTkFont(underline=True)) + self.pixel_max_preview = components.label(frame, 10, 0, pad=0, text="-", wraplength=150) + self.pixel_avg_label = components.label(frame, 9, 1, "\nAvg Pixels", pad=0, + tooltip="Average size of images in the concept by number of pixels (width * height)") + self.pixel_avg_label.configure(font=ctk.CTkFont(underline=True)) + self.pixel_avg_preview = components.label(frame, 10, 1, pad=0, text="-", wraplength=150) + self.pixel_min_label = components.label(frame, 9, 2, "\nMin Pixels", pad=0, + tooltip="Smallest image in the concept by number of pixels (width * height)") + self.pixel_min_label.configure(font=ctk.CTkFont(underline=True)) + self.pixel_min_preview = components.label(frame, 10, 2, pad=0, text="-", wraplength=150) + + #video length info + self.length_max_label = components.label(frame, 11, 0, "\nMax Length", pad=0, + tooltip="Longest video in the concept by number of frames") + self.length_max_label.configure(font=ctk.CTkFont(underline=True)) + self.length_max_preview = components.label(frame, 12, 0, pad=0, text="-", wraplength=150) + self.length_avg_label = components.label(frame, 11, 1, "\nAvg Length", pad=0, + tooltip="Average length of videos in the concept by number of frames") + self.length_avg_label.configure(font=ctk.CTkFont(underline=True)) + self.length_avg_preview = components.label(frame, 12, 1, pad=0, text="-", wraplength=150) + self.length_min_label = components.label(frame, 11, 2, "\nMin Length", pad=0, + tooltip="Shortest video in the concept by number of frames") + self.length_min_label.configure(font=ctk.CTkFont(underline=True)) + self.length_min_preview = components.label(frame, 12, 2, pad=0, text="-", wraplength=150) + + #video fps info + self.fps_max_label = components.label(frame, 13, 0, "\nMax FPS", pad=0, + tooltip="Video in concept with highest fps") + self.fps_max_label.configure(font=ctk.CTkFont(underline=True)) + self.fps_max_preview = components.label(frame, 14, 0, pad=0, text="-", wraplength=150) + self.fps_avg_label = components.label(frame, 13, 1, "\nAvg FPS", pad=0, + tooltip="Average fps of videos in the concept") + self.fps_avg_label.configure(font=ctk.CTkFont(underline=True)) + self.fps_avg_preview = components.label(frame, 14, 1, pad=0, text="-", wraplength=150) + self.fps_min_label = components.label(frame, 13, 2, "\nMin FPS", pad=0, + tooltip="Video in concept with the lowest fps") + self.fps_min_label.configure(font=ctk.CTkFont(underline=True)) + self.fps_min_preview = components.label(frame, 14, 2, pad=0, text="-", wraplength=150) + + #caption info + self.caption_max_label = components.label(frame, 15, 0, "\nMax Caption Length", pad=0, + tooltip="Largest caption in concept by character count. For token count, assume ~2 tokens/word") + self.caption_max_label.configure(font=ctk.CTkFont(underline=True)) + self.caption_max_preview = components.label(frame, 16, 0, pad=0, text="-", wraplength=150) + self.caption_avg_label = components.label(frame, 15, 1, "\nAvg Caption Length", pad=0, + tooltip="Average length of caption in concept by character count. For token count, assume ~2 tokens/word") + self.caption_avg_label.configure(font=ctk.CTkFont(underline=True)) + self.caption_avg_preview = components.label(frame, 16, 1, pad=0, text="-", wraplength=150) + self.caption_min_label = components.label(frame, 15, 2, "\nMin Caption Length", pad=0, + tooltip="Smallest caption in concept by character count. For token count, assume ~2 tokens/word") + self.caption_min_label.configure(font=ctk.CTkFont(underline=True)) + self.caption_min_preview = components.label(frame, 16, 2, pad=0, text="-", wraplength=150) + + #aspect bucket info + self.aspect_bucket_label = components.label(frame, 17, 0, "\nAspect Bucketing", pad=0, + tooltip="Graph of all possible buckets and the number of images in each one, defined as height/width. Buckets range from 0.25 (4:1 extremely wide) to 4 (1:4 extremely tall). \ + Images which don't match a bucket exactly are cropped to the nearest one.") + self.aspect_bucket_label.configure(font=ctk.CTkFont(underline=True)) + self.small_bucket_label = components.label(frame, 17, 1, "\nSmallest Buckets", pad=0, + tooltip="Image buckets with the least nonzero total images - if 'batch size' is larger than this, these images will be ignored during training! See the wiki for more details.") + self.small_bucket_label.configure(font=ctk.CTkFont(underline=True)) + self.small_bucket_preview = components.label(frame, 18, 1, pad=0, text="-") + + #aspect bucketing plot, mostly copied from timestep preview graph + appearance_mode = AppearanceModeTracker.get_mode() + background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) + text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) + background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" + self.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 + + assert self.bucket_fig is None + self.bucket_fig, self.bucket_ax = plt.subplots(figsize=(7,3)) + self.canvas = FigureCanvasTkAgg(self.bucket_fig, master=frame) + self.canvas.get_tk_widget().grid(row=19, column=0, columnspan=4, rowspan=2) + self.bucket_fig.tight_layout() + self.bucket_fig.subplots_adjust(bottom=0.15) + + self.bucket_fig.set_facecolor(background_color) + self.bucket_ax.set_facecolor(background_color) + self.bucket_ax.spines['bottom'].set_color(self.text_color) + self.bucket_ax.spines['left'].set_color(self.text_color) + self.bucket_ax.spines['top'].set_visible(False) + self.bucket_ax.spines['right'].set_color(self.text_color) + self.bucket_ax.tick_params(axis='x', colors=self.text_color, which="both") + self.bucket_ax.tick_params(axis='y', colors=self.text_color, which="both") + self.bucket_ax.xaxis.label.set_color(self.text_color) + self.bucket_ax.yaxis.label.set_color(self.text_color) + + #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") + self.refresh_advanced_stats_button = components.button(master=frame, row=0, column=1, text="Refresh Advanced", command=lambda: self.__get_concept_stats_threaded(True, 9999), + tooltip="Reload advanced statistics for the concept directory") #run "basic" scan first before "advanced", seems to help the system cache the directories and run faster + self.cancel_stats_button = components.button(master=frame, row=0, column=2, text="Abort Scan", command=lambda: self.__cancel_concept_stats(), + tooltip="Stop the currently running scan if it's taking a long time - advanced scan will be slow on large folders and on HDDs") + self.processing_time = components.label(frame, 0, 3, text="-", tooltip="Time taken to process concept directory") + + frame.pack(fill="both", expand=1) + return frame + + def __prev_image_preview(self): + self.image_preview_file_index = max(self.image_preview_file_index - 1, 0) + self.__update_image_preview() + + def __next_image_preview(self): + self.image_preview_file_index += 1 + self.__update_image_preview() + + def __update_image_preview(self): + image_preview, filename_preview, caption_preview = self.__get_preview_image() + self.image.configure(light_image=image_preview, size=image_preview.size) + self.filename_preview.configure(text=filename_preview) + self.caption_preview.configure(state="normal") + self.caption_preview.delete(index1="1.0", index2="end") + self.caption_preview.insert(index="1.0", text=caption_preview) + self.caption_preview.configure(state="disabled") + + @staticmethod + def get_concept_path(path: str) -> str | None: + if os.path.isdir(path): + return path + try: + #don't download, only check if available locally: + return huggingface_hub.snapshot_download(repo_id=path, repo_type="dataset", local_files_only=True) + except Exception: + return None + + def __download_dataset(self): + try: + huggingface_hub.login(token=self.train_config.secrets.huggingface_token, new_session=False) + huggingface_hub.snapshot_download(repo_id=self.concept.path, repo_type="dataset") + except Exception: + traceback.print_exc() + + def __download_dataset_threaded(self): + download_thread = threading.Thread(target=self.__download_dataset, daemon=True) + download_thread.start() + + def _read_text_file_for_preview(self, file_path: str) -> str: + empty_msg = "[Empty prompt]" + try: + with open(file_path, "r") as f: + if self.preview_augmentations.get(): + lines = [line.strip() for line in f if line.strip()] + return random.choice(lines) if lines else empty_msg + content = f.read().strip() + return content if content else empty_msg + except FileNotFoundError: + return "File not found, please check the path" + except IsADirectoryError: + return "[Provided path is a directory, please correct the caption path]" + except PermissionError: + if platform.system() == "Windows": + return "[Permission denied, please check the file permissions or Windows Defender settings]" + else: + return "[Permission denied, please check the file permissions]" + except UnicodeDecodeError: + return "[Invalid file encoding. This should not happen, please report this issue]" + + def __get_preview_image(self): + preview_image_path = "resources/icons/icon.png" + file_index = -1 + glob_pattern = "**/*.*" if self.concept.include_subdirectories else "*.*" + + concept_path = self.get_concept_path(self.concept.path) + if concept_path: + for path in pathlib.Path(concept_path).glob(glob_pattern): + if any(part.startswith('.') for part in path.relative_to(concept_path).parent.parts): + continue + extension = os.path.splitext(path)[1] + if path.is_file() and path_util.is_supported_image_extension(extension) \ + and not path.name.endswith("-masklabel.png") and not path.name.endswith("-condlabel.png"): + preview_image_path = path_util.canonical_join(concept_path, path) + file_index += 1 + if file_index == self.image_preview_file_index: + break + + image = load_image(preview_image_path, 'RGB') + image_tensor = functional.to_tensor(image) + + splitext = os.path.splitext(preview_image_path) + preview_mask_path = path_util.canonical_join(splitext[0] + "-masklabel.png") + if not os.path.isfile(preview_mask_path): + preview_mask_path = None + + if preview_mask_path: + mask = Image.open(preview_mask_path).convert("L") + mask_tensor = functional.to_tensor(mask) + else: + mask_tensor = torch.ones((1, image_tensor.shape[1], image_tensor.shape[2])) + + source = self.concept.text.prompt_source + preview_p = pathlib.Path(preview_image_path) + if source == "filename": + prompt_output = preview_p.stem or "[Empty prompt]" + else: + file_map = { + "sample": preview_p.with_suffix(".txt"), + "concept": pathlib.Path(self.concept.text.prompt_path) if self.concept.text.prompt_path else None, + } + file_path = file_map.get(source) + prompt_output = self._read_text_file_for_preview(str(file_path)) if file_path else "[Empty prompt]" + + modules = [] + if self.preview_augmentations.get(): + input_module = InputPipelineModule({ + 'true': True, + 'image': image_tensor, + 'mask': mask_tensor, + 'enable_random_flip': self.concept.image.enable_random_flip, + 'enable_fixed_flip': self.concept.image.enable_fixed_flip, + 'enable_random_rotate': self.concept.image.enable_random_rotate, + 'enable_fixed_rotate': self.concept.image.enable_fixed_rotate, + 'random_rotate_max_angle': self.concept.image.random_rotate_max_angle, + 'enable_random_brightness': self.concept.image.enable_random_brightness, + 'enable_fixed_brightness': self.concept.image.enable_fixed_brightness, + 'random_brightness_max_strength': self.concept.image.random_brightness_max_strength, + 'enable_random_contrast': self.concept.image.enable_random_contrast, + 'enable_fixed_contrast': self.concept.image.enable_fixed_contrast, + 'random_contrast_max_strength': self.concept.image.random_contrast_max_strength, + 'enable_random_saturation': self.concept.image.enable_random_saturation, + 'enable_fixed_saturation': self.concept.image.enable_fixed_saturation, + 'random_saturation_max_strength': self.concept.image.random_saturation_max_strength, + 'enable_random_hue': self.concept.image.enable_random_hue, + 'enable_fixed_hue': self.concept.image.enable_fixed_hue, + 'random_hue_max_strength': self.concept.image.random_hue_max_strength, + 'enable_random_circular_mask_shrink': self.concept.image.enable_random_circular_mask_shrink, + 'enable_random_mask_rotate_crop': self.concept.image.enable_random_mask_rotate_crop, + + 'prompt' : prompt_output, + 'tag_dropout_enable' : self.concept.text.tag_dropout_enable, + 'tag_dropout_probability' : self.concept.text.tag_dropout_probability, + 'tag_dropout_mode' : self.concept.text.tag_dropout_mode, + 'tag_dropout_special_tags' : self.concept.text.tag_dropout_special_tags, + 'tag_dropout_special_tags_mode' : self.concept.text.tag_dropout_special_tags_mode, + 'tag_delimiter' : self.concept.text.tag_delimiter, + 'keep_tags_count' : self.concept.text.keep_tags_count, + 'tag_dropout_special_tags_regex' : self.concept.text.tag_dropout_special_tags_regex, + 'caps_randomize_enable' : self.concept.text.caps_randomize_enable, + 'caps_randomize_probability' : self.concept.text.caps_randomize_probability, + 'caps_randomize_mode' : self.concept.text.caps_randomize_mode, + 'caps_randomize_lowercase' : self.concept.text.caps_randomize_lowercase, + 'enable_tag_shuffling' : self.concept.text.enable_tag_shuffling, + }) + + circular_mask_shrink = RandomCircularMaskShrink(mask_name='mask', shrink_probability=1.0, shrink_factor_min=0.2, shrink_factor_max=1.0, enabled_in_name='enable_random_circular_mask_shrink') + random_mask_rotate_crop = RandomMaskRotateCrop(mask_name='mask', additional_names=['image'], min_size=512, min_padding_percent=10, max_padding_percent=30, max_rotate_angle=20, enabled_in_name='enable_random_mask_rotate_crop') + random_flip = RandomFlip(names=['image', 'mask'], enabled_in_name='enable_random_flip', fixed_enabled_in_name='enable_fixed_flip') + random_rotate = RandomRotate(names=['image', 'mask'], enabled_in_name='enable_random_rotate', fixed_enabled_in_name='enable_fixed_rotate', max_angle_in_name='random_rotate_max_angle') + random_brightness = RandomBrightness(names=['image'], enabled_in_name='enable_random_brightness', fixed_enabled_in_name='enable_fixed_brightness', max_strength_in_name='random_brightness_max_strength') + random_contrast = RandomContrast(names=['image'], enabled_in_name='enable_random_contrast', fixed_enabled_in_name='enable_fixed_contrast', max_strength_in_name='random_contrast_max_strength') + random_saturation = RandomSaturation(names=['image'], enabled_in_name='enable_random_saturation', fixed_enabled_in_name='enable_fixed_saturation', max_strength_in_name='random_saturation_max_strength') + random_hue = RandomHue(names=['image'], enabled_in_name='enable_random_hue', fixed_enabled_in_name='enable_fixed_hue', max_strength_in_name='random_hue_max_strength') + drop_tags = DropTags(text_in_name='prompt', enabled_in_name='tag_dropout_enable', probability_in_name='tag_dropout_probability', dropout_mode_in_name='tag_dropout_mode', + special_tags_in_name='tag_dropout_special_tags', special_tag_mode_in_name='tag_dropout_special_tags_mode', delimiter_in_name='tag_delimiter', + keep_tags_count_in_name='keep_tags_count', text_out_name='prompt', regex_enabled_in_name='tag_dropout_special_tags_regex') + caps_randomize = CapitalizeTags(text_in_name='prompt', enabled_in_name='caps_randomize_enable', probability_in_name='caps_randomize_probability', + capitalize_mode_in_name='caps_randomize_mode', delimiter_in_name='tag_delimiter', convert_lowercase_in_name='caps_randomize_lowercase', text_out_name='prompt') + shuffle_tags = ShuffleTags(text_in_name='prompt', enabled_in_name='enable_tag_shuffling', delimiter_in_name='tag_delimiter', keep_tags_count_in_name='keep_tags_count', text_out_name='prompt') + output_module = OutputPipelineModule(['image', 'mask', 'prompt']) + + modules = [ + input_module, + circular_mask_shrink, + random_mask_rotate_crop, + random_flip, + random_rotate, + random_brightness, + random_contrast, + random_saturation, + random_hue, + drop_tags, + caps_randomize, + shuffle_tags, + output_module, + ] + + pipeline = LoadingPipeline( + device=torch.device('cpu'), + modules=modules, + batch_size=1, + seed=random.randint(0, 2**30), + state=None, + initial_epoch=0, + initial_index=0, + ) + + data = pipeline.__next__() + image_tensor = data['image'] + mask_tensor = data['mask'] + prompt_output = data['prompt'] + + filename_output = os.path.basename(preview_image_path) + + mask_tensor = torch.clamp(mask_tensor, 0.3, 1) + image_tensor = image_tensor * mask_tensor + + image = functional.to_pil_image(image_tensor) + + image.thumbnail((300, 300)) + + return image, filename_output, prompt_output + + def __update_concept_stats(self): + #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") + + #directory count + self.dir_count_preview.configure(text=self.concept.concept_stats["directory_count"]) + + #image count + self.image_count_preview.configure(text=self.concept.concept_stats["image_count"]) + self.image_count_mask_preview.configure(text=self.concept.concept_stats["image_with_mask_count"]) + self.image_count_caption_preview.configure(text=self.concept.concept_stats["image_with_caption_count"]) + + #video count + self.video_count_preview.configure(text=self.concept.concept_stats["video_count"]) + #self.video_count_mask_preview.configure(text=self.concept.concept_stats["video_with_mask_count"]) + self.video_count_caption_preview.configure(text=self.concept.concept_stats["video_with_caption_count"]) + + #mask count + self.mask_count_preview.configure(text=self.concept.concept_stats["mask_count"]) + self.mask_count_preview_unpaired.configure(text=self.concept.concept_stats["unpaired_masks"]) + + #caption count + if self.concept.concept_stats["subcaption_count"] > 0: + self.caption_count_preview.configure(text=f'{self.concept.concept_stats["caption_count"]} ({self.concept.concept_stats["subcaption_count"]})') + else: + self.caption_count_preview.configure(text=self.concept.concept_stats["caption_count"]) + self.caption_count_preview_unpaired.configure(text=self.concept.concept_stats["unpaired_captions"]) + + #resolution info + max_pixels = self.concept.concept_stats["max_pixels"] + avg_pixels = self.concept.concept_stats["avg_pixels"] + min_pixels = self.concept.concept_stats["min_pixels"] + + if any(isinstance(x, str) for x in [max_pixels, avg_pixels, min_pixels]) or self.concept.concept_stats["image_count"] == 0: #will be str if adv stats were not taken + self.pixel_max_preview.configure(text="-") + self.pixel_avg_preview.configure(text="-") + self.pixel_min_preview.configure(text="-") + else: + #formatted as (#pixels/1000000) MP, width x height, \n filename + self.pixel_max_preview.configure(text=f'{str(round(max_pixels[0]/1000000, 2))} MP, {max_pixels[2]}\n{max_pixels[1]}') + self.pixel_avg_preview.configure(text=f'{str(round(avg_pixels/1000000, 2))} MP, ~{int(math.sqrt(avg_pixels))}w x {int(math.sqrt(avg_pixels))}h') + self.pixel_min_preview.configure(text=f'{str(round(min_pixels[0]/1000000, 2))} MP, {min_pixels[2]}\n{min_pixels[1]}') + + #video length and fps info + max_length = self.concept.concept_stats["max_length"] + avg_length = self.concept.concept_stats["avg_length"] + min_length = self.concept.concept_stats["min_length"] + max_fps = self.concept.concept_stats["max_fps"] + avg_fps = self.concept.concept_stats["avg_fps"] + min_fps = self.concept.concept_stats["min_fps"] + + if any(isinstance(x, str) for x in [max_length, avg_length, min_length]) or self.concept.concept_stats["video_count"] == 0: #will be str if adv stats were not taken + self.length_max_preview.configure(text="-") + self.length_avg_preview.configure(text="-") + self.length_min_preview.configure(text="-") + self.fps_max_preview.configure(text="-") + self.fps_avg_preview.configure(text="-") + self.fps_min_preview.configure(text="-") + else: + #formatted as (#frames) frames \n filename + self.length_max_preview.configure(text=f'{int(max_length[0])} frames\n{max_length[1]}') + self.length_avg_preview.configure(text=f'{int(avg_length)} frames') + self.length_min_preview.configure(text=f'{int(min_length[0])} frames\n{min_length[1]}') + #formatted as (#fps) fps \n filename + self.fps_max_preview.configure(text=f'{int(max_fps[0])} fps\n{max_fps[1]}') + self.fps_avg_preview.configure(text=f'{int(avg_fps)} fps') + self.fps_min_preview.configure(text=f'{int(min_fps[0])} fps\n{min_fps[1]}') + + #caption info + max_caption_length = self.concept.concept_stats["max_caption_length"] + avg_caption_length = self.concept.concept_stats["avg_caption_length"] + min_caption_length = self.concept.concept_stats["min_caption_length"] + + if any(isinstance(x, str) for x in [max_caption_length, avg_caption_length, min_caption_length]) or self.concept.concept_stats["caption_count"] == 0: #will be str if adv stats were not taken + self.caption_max_preview.configure(text="-") + self.caption_avg_preview.configure(text="-") + self.caption_min_preview.configure(text="-") + else: + #formatted as (#chars) chars, (#words) words, \n filename + self.caption_max_preview.configure(text=f'{max_caption_length[0]} chars, {max_caption_length[2]} words\n{max_caption_length[1]}') + self.caption_avg_preview.configure(text=f'{int(avg_caption_length[0])} chars, {int(avg_caption_length[1])} words') + self.caption_min_preview.configure(text=f'{min_caption_length[0]} chars, {min_caption_length[2]} words\n{min_caption_length[1]}') + + #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) + + self.bucket_ax.cla() + aspects = [str(x) for x in list(aspect_buckets.keys())] + aspect_ratios = [self.decimal_to_aspect_ratio(x) for x in list(aspect_buckets.keys())] + counts = list(aspect_buckets.values()) + b = self.bucket_ax.bar(aspect_ratios, counts) + self.bucket_ax.bar_label(b, color=self.text_color) + sec = self.bucket_ax.secondary_xaxis(location=-0.1) + sec.spines["bottom"].set_linewidth(0) + sec.set_xticks([0, (len(aspects)-1)/2, len(aspects)-1], labels=["Wide", "Square", "Tall"]) + sec.tick_params('x', length=0) + self.canvas.draw() + + def decimal_to_aspect_ratio(self, value : float): + #find closest fraction to decimal aspect value and convert to a:b format + aspect_fraction = fractions.Fraction(value).limit_denominator(16) + aspect_string = f'{aspect_fraction.denominator}:{aspect_fraction.numerator}' + return aspect_string + + def __get_concept_stats(self, advanced_checks: bool, wait_time: float): + start_time = time.perf_counter() + last_update = time.perf_counter() + self.cancel_scan_flag.clear() + self.concept_stats_tab.after(0, self.__disable_scan_buttons) + concept_path = self.get_concept_path(self.concept.path) + + if not concept_path: + print(f"Unable to get statistics for concept path: {self.concept.path}") + self.concept_stats_tab.after(0, self.__enable_scan_buttons) + return + subfolders = [concept_path] + + 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: + 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 + subfolders.extend([f for f in os.scandir(path) if f.is_dir() and not f.name.startswith('.')]) + self.concept.concept_stats = stats_dict + #update GUI approx every half second + if time.perf_counter() > (last_update + 0.5): + 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) + + 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): + self.refresh_basic_stats_button.configure(state="disabled") + self.refresh_advanced_stats_button.configure(state="disabled") + + def __enable_scan_buttons(self): + self.refresh_basic_stats_button.configure(state="normal") + self.refresh_advanced_stats_button.configure(state="normal") + + def __cancel_concept_stats(self): + self.cancel_scan_flag.set() + + def __auto_update_concept_stats(self): + try: + self.__update_concept_stats() #load stats from config if available, else raises KeyError + if self.concept.concept_stats["file_size"] == 0: #force rescan if empty + raise KeyError + except KeyError: + concept_path = self.get_concept_path(self.concept.path) + if concept_path: + self.__get_concept_stats(False, 2) #force rescan if config is empty, timeout of 2 sec + if self.concept.concept_stats["processing_time"] < 0.1: + self.__get_concept_stats(True, 2) #do advanced scan automatically if basic took <0.1s + + def destroy(self): + if self.bucket_fig is not None: + plt.close(self.bucket_fig) + self.bucket_fig = None + + super().destroy() + + def __ok(self): + self.destroy() diff --git a/modules/ui/ConvertModelUIController.py b/modules/ui/ConvertModelUIController.py new file mode 100644 index 000000000..6cb1b507a --- /dev/null +++ b/modules/ui/ConvertModelUIController.py @@ -0,0 +1,170 @@ +import traceback +from uuid import uuid4 + +from modules.util import create +from modules.util.args.ConvertModelArgs import ConvertModelArgs +from modules.util.config.TrainConfig import QuantizationConfig +from modules.util.enum.DataType import DataType +from modules.util.enum.ModelFormat import ModelFormat +from modules.util.enum.ModelType import ModelType +from modules.util.enum.PathIOType import PathIOType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.ModelNames import EmbeddingName, ModelNames +from modules.util.torch_util import torch_gc +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class ConvertModelUI(ctk.CTkToplevel): + def __init__(self, parent, *args, **kwargs): + super().__init__(parent, *args, **kwargs) + self.parent = parent + + self.parent = parent + self.convert_model_args = ConvertModelArgs.default_values() + self.ui_state = UIState(self, self.convert_model_args) + self.button = None + + + self.title("Convert models") + self.geometry("550x350") + self.resizable(True, True) + + self.frame = ctk.CTkFrame(self, width=600, height=300) + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + self.frame.grid_columnconfigure(0, weight=0) + self.frame.grid_columnconfigure(1, weight=1) + + self.main_frame(self.frame) + self.frame.pack(fill="both", expand=True) + + self.wait_visibility() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def main_frame(self, master): + # model type + components.label(master, 0, 0, "Model Type", + tooltip="Type of the model") + components.options_kv(master, 0, 1, [ #TODO simplify + ("Stable Diffusion 1.5", ModelType.STABLE_DIFFUSION_15), + ("Stable Diffusion 1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), + ("Stable Diffusion 2.0", ModelType.STABLE_DIFFUSION_20), + ("Stable Diffusion 2.0 Inpainting", ModelType.STABLE_DIFFUSION_20_INPAINTING), + ("Stable Diffusion 2.1", ModelType.STABLE_DIFFUSION_21), + ("Stable Diffusion 3", ModelType.STABLE_DIFFUSION_3), + ("Stable Diffusion 3.5", ModelType.STABLE_DIFFUSION_35), + ("Stable Diffusion XL 1.0 Base", ModelType.STABLE_DIFFUSION_XL_10_BASE), + ("Stable Diffusion XL 1.0 Base Inpainting", ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING), + ("Wuerstchen v2", ModelType.WUERSTCHEN_2), + ("Stable Cascade", ModelType.STABLE_CASCADE_1), + ("PixArt Alpha", ModelType.PIXART_ALPHA), + ("PixArt Sigma", ModelType.PIXART_SIGMA), + ("Flux Dev", ModelType.FLUX_DEV_1), + ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), + ("Flux 2", ModelType.FLUX_2), + ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), + ("Chroma1", ModelType.CHROMA_1), #TODO does this just work? HiDream is not here + ("QwenImage", ModelType.QWEN), #TODO does this just work? HiDream is not here + ("ZImage", ModelType.Z_IMAGE), + ], self.ui_state, "model_type") + + # training method + components.label(master, 1, 0, "Model Type", + tooltip="The type of model to convert") + components.options_kv(master, 1, 1, [ + ("Base Model", TrainingMethod.FINE_TUNE), + ("LoRA", TrainingMethod.LORA), + ("Embedding", TrainingMethod.EMBEDDING), + ], self.ui_state, "training_method") + + # input name + components.label(master, 2, 0, "Input name", + tooltip="Filename, directory or hugging face repository of the base model") + components.path_entry( + master, 2, 1, self.ui_state, "input_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # output data type + components.label(master, 3, 0, "Output Data Type", + tooltip="Precision to use when saving the output model") + components.options_kv(master, 3, 1, [ + ("float32", DataType.FLOAT_32), + ("float16", DataType.FLOAT_16), + ("bfloat16", DataType.BFLOAT_16), + ], self.ui_state, "output_dtype") + + # output format + components.label(master, 4, 0, "Output Format", + tooltip="Format to use when saving the output model") + components.options_kv(master, 4, 1, [ + ("Safetensors", ModelFormat.SAFETENSORS), + ("Diffusers", ModelFormat.DIFFUSERS), + ], self.ui_state, "output_model_format") + + # output model destination + components.label(master, 5, 0, "Model Output Destination", + tooltip="Filename or directory where the output model is saved") + components.path_entry( + master, 5, 1, self.ui_state, "output_model_destination", + mode="file", + io_type=PathIOType.MODEL, + ) + + self.button = components.button(master, 6, 1, "Convert", self.convert_model) + + def convert_model(self): + try: + self.button.configure(state="disabled") + model_loader = create.create_model_loader( + model_type=self.convert_model_args.model_type, + training_method=self.convert_model_args.training_method + ) + model_saver = create.create_model_saver( + model_type=self.convert_model_args.model_type, + training_method=self.convert_model_args.training_method + ) + + print("Loading model " + self.convert_model_args.input_name) + if self.convert_model_args.training_method in [TrainingMethod.FINE_TUNE]: + model = model_loader.load( + model_type=self.convert_model_args.model_type, + model_names=ModelNames( + base_model=self.convert_model_args.input_name, + ), + weight_dtypes=self.convert_model_args.weight_dtypes(), + quantization=QuantizationConfig.default_values(), + ) + elif self.convert_model_args.training_method in [TrainingMethod.LORA, TrainingMethod.EMBEDDING]: + model = model_loader.load( + model_type=self.convert_model_args.model_type, + model_names=ModelNames( + base_model=None, + lora=self.convert_model_args.input_name, + embedding=EmbeddingName(str(uuid4()), self.convert_model_args.input_name), + ), + weight_dtypes=self.convert_model_args.weight_dtypes(), + quantization=QuantizationConfig.default_values(), + ) + else: + raise Exception("could not load model: " + self.convert_model_args.input_name) + + print("Saving model " + self.convert_model_args.output_model_destination) + model_saver.save( + model=model, + model_type=self.convert_model_args.model_type, + output_model_format=self.convert_model_args.output_model_format, + output_model_destination=self.convert_model_args.output_model_destination, + dtype=self.convert_model_args.output_dtype.torch_dtype(), + ) + print("Model converted") + except Exception: + traceback.print_exc() + + torch_gc() + self.button.configure(state="normal") diff --git a/modules/ui/CtkAdditionalEmbeddingsTabView.py b/modules/ui/CtkAdditionalEmbeddingsTabView.py new file mode 100644 index 000000000..6a5e3fbe7 --- /dev/null +++ b/modules/ui/CtkAdditionalEmbeddingsTabView.py @@ -0,0 +1,136 @@ + +from modules.ui.ConfigList import ConfigList +from modules.util.config.TrainConfig import TrainConfig, TrainEmbeddingConfig +from modules.util.ui import components +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class AdditionalEmbeddingsTab(ConfigList): + + def __init__(self, master, train_config: TrainConfig, ui_state: UIState): + super().__init__( + master, + train_config, + ui_state, + attr_name="additional_embeddings", + enable_key="train", + from_external_file=False, + add_button_text="add embedding", + is_full_width=True, + show_toggle_button=True + ) + + def refresh_ui(self): + if self.element_list is not None: + self.element_list.destroy() + self.element_list = None + self.widgets_initialized = False + self._create_element_list() + + def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): + return EmbeddingWidget(master, element, i, open_command, remove_command, clone_command, save_command) + + def create_new_element(self) -> dict: + return TrainEmbeddingConfig.default_values() + + def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: + pass + + +class EmbeddingWidget(ctk.CTkFrame): + def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): + super().__init__( + master=master, corner_radius=10, bg_color="transparent" + ) + + self.element = element + self.ui_state = UIState(self, element) + self.i = i + self.save_command = save_command + + self.grid_columnconfigure(0, weight=1) + + top_frame = ctk.CTkFrame(master=self, corner_radius=0, fg_color="transparent") + top_frame.grid(row=0, column=0, sticky="nsew") + top_frame.grid_columnconfigure(3, weight=1) + top_frame.grid_columnconfigure(5, weight=1) + + bottom_frame = ctk.CTkFrame(master=self, corner_radius=0, fg_color="transparent") + bottom_frame.grid(row=1, column=0, sticky="nsew") + bottom_frame.grid_columnconfigure(7, weight=1) + + # close button + close_button = ctk.CTkButton( + master=top_frame, + width=20, + height=20, + text="X", + corner_radius=2, + fg_color="#C00000", + command=lambda: remove_command(self.i), + ) + close_button.grid(row=0, column=0) + + # clone button + clone_button = ctk.CTkButton( + master=top_frame, + width=20, + height=20, + text="+", + corner_radius=2, + fg_color="#00C000", + command=lambda: clone_command(self.i, self.__randomize_uuid), + ) + clone_button.grid(row=0, column=1, padx=5) + + # 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.path_entry( + top_frame, 0, 3, self.ui_state, "model_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # placeholder + components.label(top_frame, 0, 4, "placeholder:", + tooltip="The placeholder used when using the embedding in a prompt") + components.entry(top_frame, 0, 5, self.ui_state, "placeholder") + + # token count + components.label(top_frame, 0, 6, "token count:", + tooltip="The token count used when creating a new embedding. Leave empty to auto detect from the initial embedding text.") + token_count_entry = components.entry(top_frame, 0, 7, self.ui_state, "token_count") + token_count_entry.configure(width=40) + + # trainable + components.label(bottom_frame, 0, 0, "train:") + trainable_switch = components.switch(bottom_frame, 0, 1, self.ui_state, "train", command=save_command) + trainable_switch.configure(width=40) + + # output embedding + components.label(bottom_frame, 0, 2, "output embedding:", + tooltip="Output embeddings are calculated at the output of the text encoder, not the input. This can improve results for larger text encoders and lower VRAM usage.") + output_embedding_switch = components.switch(bottom_frame, 0, 3, self.ui_state, "is_output_embedding") + output_embedding_switch.configure(width=40) + + # stop training after + components.label(bottom_frame, 0, 4, "stop training after:", + tooltip="When to stop training the embedding") + components.time_entry(bottom_frame, 0, 5, self.ui_state, "stop_training_after", "stop_training_after_unit") + + # initial embedding text + components.label(bottom_frame, 0, 6, "initial embedding text:", + tooltip="The initial embedding text used when creating a new embedding") + components.entry(bottom_frame, 0, 7, self.ui_state, "initial_embedding_text") + + def __randomize_uuid(self, embedding_config: TrainEmbeddingConfig): + embedding_config.uuid = TrainEmbeddingConfig.default_values().uuid + return embedding_config + + def configure_element(self): + pass + + def place_in_list(self): + self.grid(row=self.i, column=0, pady=5, padx=5, sticky="new") diff --git a/modules/ui/CtkCaptionUIView.py b/modules/ui/CtkCaptionUIView.py new file mode 100644 index 000000000..e6cc0551e --- /dev/null +++ b/modules/ui/CtkCaptionUIView.py @@ -0,0 +1,572 @@ +import os +import platform +import subprocess +import traceback +from tkinter import filedialog + +from modules.module.Blip2Model import Blip2Model +from modules.module.BlipModel import BlipModel +from modules.module.ClipSegModel import ClipSegModel +from modules.module.MaskByColor import MaskByColor +from modules.module.RembgHumanModel import RembgHumanModel +from modules.module.RembgModel import RembgModel +from modules.module.WDModel import WDModel +from modules.ui.GenerateCaptionsWindow import GenerateCaptionsWindow +from modules.ui.GenerateMasksWindow import GenerateMasksWindow +from modules.util import path_util +from modules.util.image_util import load_image +from modules.util.torch_util import default_device, torch_gc +from modules.util.ui import components +from modules.util.ui.ui_utils import bind_mousewheel, set_window_icon +from modules.util.ui.UIState import UIState + +import torch + +import customtkinter as ctk +import cv2 +import numpy as np +from customtkinter import ScalingTracker, ThemeManager +from PIL import Image, ImageDraw + + +class CaptionUI(ctk.CTkToplevel): + def __init__( + self, + parent, + initial_dir: str | None, + initial_include_subdirectories: bool, + *args, + **kwargs, + ) -> None: + super().__init__(parent, *args, **kwargs) + self.protocol("WM_DELETE_WINDOW", self._on_close) + + self.dir = initial_dir + self.config_ui_data = {"include_subdirectories": initial_include_subdirectories} + self.config_ui_state = UIState(self, self.config_ui_data) + self.image_size = 850 + self.help_text = """ + Keyboard shortcuts when focusing on the prompt input field: + Up arrow: previous image + Down arrow: next image + Return: save + Ctrl+M: only show the mask + Ctrl+D: draw mask editing mode + Ctrl+F: fill mask editing mode + + When editing masks: + Left click: add mask + Right click: remove mask + Mouse wheel: increase or decrease brush size""" + self.masking_model = None + self.captioning_model = None + self.image_rel_paths = [] + self.current_image_index = -1 + self.file_list = None + self.image_labels = [] + self.pil_image = None + self.image_width = 0 + self.image_height = 0 + self.pil_mask = None + self.mask_draw_x = 0 + self.mask_draw_y = 0 + self.mask_draw_radius = 0.01 + self.display_only_mask = False + self.image = None + self.image_label = None + self.mask_editing_mode = 'draw' + self.enable_mask_editing_var = ctk.BooleanVar() + self.mask_editing_alpha = None + self.prompt_var = None + self.prompt_component = None + + + self.title("OneTrainer") + self.geometry("1280x980") + self.resizable(False, False) + + + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=1) + self.grid_columnconfigure(0, weight=1) + + + self.top_bar(self) + + self.bottom_frame = ctk.CTkFrame(self) + self.bottom_frame.grid(row=1, column=0, sticky="nsew") + self.bottom_frame.grid_rowconfigure(0, weight=1) + self.bottom_frame.grid_columnconfigure(0, weight=0) + self.bottom_frame.grid_columnconfigure(1, weight=1) + + self.file_list_column(self.bottom_frame) + self.content_column(self.bottom_frame) + self.load_directory() + + self.wait_visibility() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + def top_bar(self, master): + top_frame = ctk.CTkFrame(master) + top_frame.grid(row=0, column=0, sticky="nsew") + + components.button(top_frame, 0, 0, "Open", self.open_directory, + tooltip="open a new directory") + components.button(top_frame, 0, 1, "Generate Masks", self.open_mask_window, + tooltip="open a dialog to automatically generate masks") + components.button(top_frame, 0, 2, "Generate Captions", self.open_caption_window, + tooltip="open a dialog to automatically generate captions") + + if platform.system() == "Windows": + components.button(top_frame, 0, 3, "Open in Explorer", self.open_in_explorer, + tooltip="open the current image in Explorer") + + components.switch(top_frame, 0, 4, self.config_ui_state, "include_subdirectories", + text="include subdirectories") + + top_frame.grid_columnconfigure(5, weight=1) + + components.button(top_frame, 0, 6, "Help", self.print_help, + tooltip=self.help_text) + + def file_list_column(self, master): + if self.file_list is not None: + self.image_labels = [] + self.file_list.destroy() + + self.file_list = ctk.CTkScrollableFrame(master, width=300) + self.file_list.grid(row=0, column=0, sticky="nsew") + + for i, filename in enumerate(self.image_rel_paths): + def __create_switch_image(index): + def __switch_image(event): + self.switch_image(index) + + return __switch_image + + label = ctk.CTkLabel(self.file_list, text=filename) + label.bind("", __create_switch_image(i)) + + self.image_labels.append(label) + label.grid(row=i, column=0, padx=5, sticky="nsw") + + def content_column(self, master): + image = Image.new("RGBA", (512, 512), (0, 0, 0, 0)) + + right_frame = ctk.CTkFrame(master, fg_color="transparent") + right_frame.grid(row=0, column=1, sticky="nsew") + + right_frame.grid_columnconfigure(4, weight=1) + right_frame.grid_rowconfigure(1, weight=1) + + components.button(right_frame, 0, 0, "Draw", self.draw_mask_editing_mode, + tooltip="draw a mask using a brush") + components.button(right_frame, 0, 1, "Fill", self.fill_mask_editing_mode, + tooltip="draw a mask using a fill tool") + + # checkbox to enable mask editing + self.enable_mask_editing_var = ctk.BooleanVar() + self.enable_mask_editing_var.set(False) + enable_mask_editing_checkbox = ctk.CTkCheckBox( + right_frame, text="Enable Mask Editing", variable=self.enable_mask_editing_var, width=50) + enable_mask_editing_checkbox.grid(row=0, column=2, padx=25, pady=5, sticky="w") + + # mask alpha textbox + self.mask_editing_alpha = ctk.CTkEntry(master=right_frame, width=40, placeholder_text="1.0") + self.mask_editing_alpha.insert(0, "1.0") + self.mask_editing_alpha.grid(row=0, column=3, sticky="e", padx=5, pady=5) + self.bind_key_events(self.mask_editing_alpha) + + mask_editing_alpha_label = ctk.CTkLabel(right_frame, text="Brush Alpha", width=75) + mask_editing_alpha_label.grid(row=0, column=4, padx=0, pady=5, sticky="w") + + # image + self.image = ctk.CTkImage( + light_image=image, + size=(self.image_size, self.image_size) + ) + self.image_label = ctk.CTkLabel( + master=right_frame, text="", image=self.image, height=self.image_size, width=self.image_size + ) + self.image_label.grid(row=1, column=0, columnspan=5, sticky="nsew") + + self.image_label.bind("", self.edit_mask) + self.image_label.bind("", self.edit_mask) + self.image_label.bind("", self.edit_mask) + bind_mousewheel(self.image_label, {self.image_label.children["!label"]}, self.draw_mask_radius) + + # prompt + self.prompt_var = ctk.StringVar() + self.prompt_component = ctk.CTkEntry(right_frame, textvariable=self.prompt_var) + self.prompt_component.grid(row=2, column=0, columnspan=5, pady=5, sticky="new") + self.bind_key_events(self.prompt_component) + self.prompt_component.focus_set() + + def bind_key_events(self, component): + component.bind("", self.next_image) + component.bind("", self.previous_image) + component.bind("", self.save) + component.bind("", self.toggle_mask) + component.bind("", self.draw_mask_editing_mode) + component.bind("", self.fill_mask_editing_mode) + + def load_directory(self, include_subdirectories: bool = False): + self.scan_directory(include_subdirectories) + self.file_list_column(self.bottom_frame) + + if len(self.image_rel_paths) > 0: + self.switch_image(0) + else: + self.switch_image(-1) + + self.prompt_component.focus_set() + + def scan_directory(self, include_subdirectories: bool = False): + def __is_supported_image_extension(filename): + name, ext = os.path.splitext(filename) + return path_util.is_supported_image_extension(ext) and not name.endswith("-masklabel") and not name.endswith("-condlabel") + + self.image_rel_paths = [] + + if not self.dir or not os.path.isdir(self.dir): + return + + if include_subdirectories: + for root, _, files in os.walk(self.dir): + for filename in files: + if __is_supported_image_extension(filename): + self.image_rel_paths.append( + os.path.relpath(os.path.join(root, filename), self.dir) + ) + else: + for _, filename in enumerate(os.listdir(self.dir)): + if __is_supported_image_extension(filename): + self.image_rel_paths.append( + os.path.relpath(os.path.join(self.dir, filename), self.dir) + ) + + def load_image(self): + image_name = "resources/icons/icon.png" + + if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): + image_name = self.image_rel_paths[self.current_image_index] + image_name = os.path.join(self.dir, image_name) + + try: + return load_image(image_name, convert_mode="RGB") + except Exception: + print(f'Could not open image {image_name}') + + def load_mask(self): + if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): + image_name = self.image_rel_paths[self.current_image_index] + mask_name = os.path.splitext(image_name)[0] + "-masklabel.png" + mask_name = os.path.join(self.dir, mask_name) + + try: + return load_image(mask_name, convert_mode='RGB') + except Exception: + return None + else: + return None + + def load_prompt(self): + if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): + image_name = self.image_rel_paths[self.current_image_index] + prompt_name = os.path.splitext(image_name)[0] + ".txt" + prompt_name = os.path.join(self.dir, prompt_name) + + try: + with open(prompt_name, "r", encoding='utf-8') as f: + return f.readlines()[0].strip() + except Exception: + return "" + else: + return "" + + def previous_image(self, event): + if len(self.image_rel_paths) > 0 and (self.current_image_index - 1) >= 0: + self.switch_image(self.current_image_index - 1) + + def next_image(self, event): + if len(self.image_rel_paths) > 0 and (self.current_image_index + 1) < len(self.image_rel_paths): + self.switch_image(self.current_image_index + 1) + + def switch_image(self, index): + if len(self.image_labels) > 0 and self.current_image_index < len(self.image_labels): + self.image_labels[self.current_image_index].configure( + text_color=ThemeManager.theme["CTkLabel"]["text_color"]) + + self.current_image_index = index + if index >= 0: + self.image_labels[index].configure(text_color="#FF0000") + + self.pil_image = self.load_image() + self.pil_mask = self.load_mask() + prompt = self.load_prompt() + + self.image_width = self.pil_image.width + self.image_height = self.pil_image.height + scale = self.image_size / max(self.pil_image.height, self.pil_image.width) + height = int(self.pil_image.height * scale) + width = int(self.pil_image.width * scale) + + self.pil_image = self.pil_image.resize((width, height), Image.Resampling.LANCZOS) + + self.refresh_image() + self.prompt_var.set(prompt) + else: + image = Image.new("RGB", (512, 512), (0, 0, 0)) + self.image.configure(light_image=image) + + def refresh_image(self): + if self.pil_mask: + resized_pil_mask = self.pil_mask.resize( + (self.pil_image.width, self.pil_image.height), + Image.Resampling.NEAREST + ) + + if self.display_only_mask: + self.image.configure(light_image=resized_pil_mask, size=resized_pil_mask.size) + else: + np_image = np.array(self.pil_image).astype(np.float32) / 255.0 + np_mask = np.array(resized_pil_mask).astype(np.float32) / 255.0 + + # normalize mask between 0.3 - 1.0 so we can see image underneath and gauge strength of the alpha + norm_min = 0.3 + np_mask_min = np_mask.min() + if np_mask_min == 0: + # optimize for common case + np_mask = np_mask * (1.0 - norm_min) + norm_min + elif np_mask_min < 1: + # note: min of 1 means we get divide by 0 + np_mask = (np_mask - np_mask_min) / (1.0 - np_mask_min) * (1.0 - norm_min) + norm_min + + np_masked_image = (np_image * np_mask * 255.0).astype(np.uint8) + masked_image = Image.fromarray(np_masked_image, mode='RGB') + + self.image.configure(light_image=masked_image, size=masked_image.size) + else: + self.image.configure(light_image=self.pil_image, size=self.pil_image.size) + + def draw_mask_radius(self, delta, raw_event): + # Wheel up = Increase radius. Wheel down = Decrease radius. + multiplier = 1.0 + (delta * 0.05) + self.mask_draw_radius = max(0.0025, self.mask_draw_radius * multiplier) + + def edit_mask(self, event): + if not self.enable_mask_editing_var.get(): + return + + if event.widget != self.image_label.children["!label"]: + return + + if len(self.image_rel_paths) == 0 or self.current_image_index >= len(self.image_rel_paths): + return + + display_scaling = ScalingTracker.get_window_scaling(self) + + event_x = event.x / display_scaling + event_y = event.y / display_scaling + + start_x = int(event_x / self.pil_image.width * self.image_width) + start_y = int(event_y / self.pil_image.height * self.image_height) + end_x = int(self.mask_draw_x / self.pil_image.width * self.image_width) + end_y = int(self.mask_draw_y / self.pil_image.height * self.image_height) + + self.mask_draw_x = event_x + self.mask_draw_y = event_y + + is_right = False + is_left = False + if event.state & 0x0100 or event.num == 1: # left mouse button + is_left = True + elif event.state & 0x0400 or event.num == 3: # right mouse button + is_right = True + + if self.mask_editing_mode == 'draw': + self.draw_mask(start_x, start_y, end_x, end_y, is_left, is_right) + if self.mask_editing_mode == 'fill': + self.fill_mask(start_x, start_y, end_x, end_y, is_left, is_right) + + def draw_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): + color = None + + adding_to_mask = True + if is_left: + try: + alpha = float(self.mask_editing_alpha.get()) + except Exception: + alpha = 1.0 + rgb_value = int(max(0.0, min(alpha, 1.0)) * 255) # max/min stuff to clamp to 0 - 255 range + color = (rgb_value, rgb_value, rgb_value) + + elif is_right: + color = (0, 0, 0) + adding_to_mask = False + + if color is not None: + if self.pil_mask is None: + if adding_to_mask: + self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(0, 0, 0)) + else: + self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(255, 255, 255)) + + radius = int(self.mask_draw_radius * max(self.pil_mask.width, self.pil_mask.height)) + + draw = ImageDraw.Draw(self.pil_mask) + draw.line((start_x, start_y, end_x, end_y), fill=color, + width=radius + radius + 1) + draw.ellipse((start_x - radius, start_y - radius, + start_x + radius, start_y + radius), fill=color, outline=None) + draw.ellipse((end_x - radius, end_y - radius, end_x + radius, + end_y + radius), fill=color, outline=None) + + self.refresh_image() + + def fill_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): + color = None + + adding_to_mask = True + if is_left: + try: + alpha = float(self.mask_editing_alpha.get()) + except Exception: + alpha = 1.0 + rgb_value = int(max(0.0, min(alpha, 1.0)) * 255) # max/min stuff to clamp to 0 - 255 range + color = (rgb_value, rgb_value, rgb_value) + + elif is_right: + color = (0, 0, 0) + adding_to_mask = False + + if color is not None: + if self.pil_mask is None: + if adding_to_mask: + self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(0, 0, 0)) + else: + self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(255, 255, 255)) + + np_mask = np.array(self.pil_mask).astype(np.uint8) + cv2.floodFill(np_mask, None, (start_x, start_y), color) + self.pil_mask = Image.fromarray(np_mask, 'RGB') + + self.refresh_image() + + def save(self, event): + if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): + image_name = self.image_rel_paths[self.current_image_index] + + prompt_name = os.path.splitext(image_name)[0] + ".txt" + prompt_name = os.path.join(self.dir, prompt_name) + + mask_name = os.path.splitext(image_name)[0] + "-masklabel.png" + mask_name = os.path.join(self.dir, mask_name) + + try: + with open(prompt_name, "w", encoding='utf-8') as f: + f.write(self.prompt_var.get()) + except Exception: + return + + if self.pil_mask: + self.pil_mask.save(mask_name) + + def draw_mask_editing_mode(self, *args): + self.mask_editing_mode = 'draw' + + if args: + # disable default event + return "break" + return None + + def fill_mask_editing_mode(self, *args): + self.mask_editing_mode = 'fill' + + def toggle_mask(self, *args): + self.display_only_mask = not self.display_only_mask + self.refresh_image() + + def open_directory(self): + new_dir = filedialog.askdirectory() + + if new_dir: + self.dir = new_dir + self.load_directory(include_subdirectories=self.config_ui_data["include_subdirectories"]) + + def open_mask_window(self): + dialog = GenerateMasksWindow(self, self.dir, self.config_ui_data["include_subdirectories"]) + self.wait_window(dialog) + self.switch_image(self.current_image_index) + + def open_caption_window(self): + dialog = GenerateCaptionsWindow(self, self.dir, self.config_ui_data["include_subdirectories"]) + self.wait_window(dialog) + self.switch_image(self.current_image_index) + + def open_in_explorer(self): + try: + image_name = self.image_rel_paths[self.current_image_index] + image_name = os.path.realpath(os.path.join(self.dir, image_name)) + subprocess.Popen(f"explorer /select,{image_name}") + except Exception: + traceback.print_exc() + + def load_masking_model(self, model): + model_type = type(self.masking_model).__name__ if self.masking_model else None + + if model == "ClipSeg" and model_type != "ClipSegModel": + self._release_models() + print("loading ClipSeg model, this may take a while") + self.masking_model = ClipSegModel(default_device, torch.float32) + elif model == "Rembg" and model_type != "RembgModel": + self._release_models() + print("loading Rembg model, this may take a while") + self.masking_model = RembgModel(default_device, torch.float32) + elif model == "Rembg-Human" and model_type != "RembgHumanModel": + self._release_models() + print("loading Rembg-Human model, this may take a while") + self.masking_model = RembgHumanModel(default_device, torch.float32) + elif model == "Hex Color" and model_type != "MaskByColor": + self._release_models() + self.masking_model = MaskByColor(default_device, torch.float32) + + def load_captioning_model(self, model): + model_type = type(self.captioning_model).__name__ if self.captioning_model else None + + if model == "Blip" and model_type != "BlipModel": + self._release_models() + print("loading Blip model, this may take a while") + self.captioning_model = BlipModel(default_device, torch.float16) + elif model == "Blip2" and model_type != "Blip2Model": + self._release_models() + print("loading Blip2 model, this may take a while") + self.captioning_model = Blip2Model(default_device, torch.float16) + elif model == "WD14 VIT v2" and model_type != "WDModel": + self._release_models() + print("loading WD14_VIT_v2 model, this may take a while") + self.captioning_model = WDModel(default_device, torch.float16) + + def print_help(self): + print(self.help_text) + + def _release_models(self): + """Release all models from VRAM""" + freed = False + if self.captioning_model is not None: + self.captioning_model = None + freed = True + if self.masking_model is not None: + self.masking_model = None + freed = True + if freed: + torch_gc() + + def _on_close(self): + self._release_models() + self.destroy() + + def destroy(self): + self._release_models() + super().destroy() diff --git a/modules/ui/CtkCloudTabView.py b/modules/ui/CtkCloudTabView.py new file mode 100644 index 000000000..99057e428 --- /dev/null +++ b/modules/ui/CtkCloudTabView.py @@ -0,0 +1,221 @@ + +import webbrowser + +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.CloudAction import CloudAction +from modules.util.enum.CloudFileSync import CloudFileSync +from modules.util.enum.CloudType import CloudType +from modules.util.ui import components +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class CloudTab: + + def __init__(self, master, train_config: TrainConfig, ui_state: UIState, parent): + super().__init__() + + self.master = master + self.train_config = train_config + self.ui_state = ui_state + self.parent = parent + self.reattach = False + + self.frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + self.frame.grid_columnconfigure(0, weight=0) + self.frame.grid_columnconfigure(1, weight=1) + self.frame.grid_columnconfigure(2, weight=0) + self.frame.grid_columnconfigure(3, weight=1) + self.frame.grid_columnconfigure(4, weight=0) + self.frame.grid_columnconfigure(5, weight=1) + + components.label(self.frame, 0, 0, "Enabled", + tooltip="Enable cloud training") + components.switch(self.frame, 0, 1, self.ui_state, "cloud.enabled") + + components.label(self.frame, 1, 0, "Type", + tooltip="Choose LINUX to connect to a linux machine via SSH. Choose RUNPOD for additional functionality such as automatically creating and deleting pods.") + components.options_kv(self.frame, 1, 1, [ + ("RUNPOD", CloudType.RUNPOD), + ("LINUX", CloudType.LINUX), + ], self.ui_state, "cloud.type") + + components.label(self.frame, 2, 0, "File sync method", + tooltip="Choose NATIVE_SCP to use scp.exe to transfer files. FABRIC_SFTP uses the Paramiko/Fabric SFTP implementation for file transfers instead.") + components.options_kv(self.frame, 2, 1, [ + ("NATIVE_SCP", CloudFileSync.NATIVE_SCP), + ("FABRIC_SFTP", CloudFileSync.FABRIC_SFTP), + ], self.ui_state, "cloud.file_sync") + + components.label(self.frame, 3, 0, "API key", + tooltip="Cloud service API key for RUNPOD. Leave empty for LINUX. This value is stored separately, not saved to your configuration file. ") + components.entry(self.frame, 3, 1, self.ui_state, "secrets.cloud.api_key") + + components.label(self.frame, 4, 0, "Hostname", + tooltip="SSH server hostname or IP. Leave empty if you have a Cloud ID or want to automatically create a new cloud.") + components.entry(self.frame, 4, 1, self.ui_state, "secrets.cloud.host") + + components.label(self.frame, 5, 0, "Port", + tooltip="SSH server port. Leave empty if you have a Cloud ID or want to automatically create a new cloud.") + components.entry(self.frame, 5, 1, self.ui_state, "secrets.cloud.port") + + components.label(self.frame, 6, 0, "User", + tooltip='SSH username. Use "root" for RUNPOD. Your SSH client must be set up to connect to the cloud using a public key, without a password. For RUNPOD, create an ed25519 key locally, and copy the contents of the public keyfile to your "SSH Public Keys" on the RunPod website.') + components.entry(self.frame, 6, 1, self.ui_state, "secrets.cloud.user") + + 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.path_entry(self.frame, 7, 1, self.ui_state, "secrets.cloud.key_file", mode="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.") + components.entry(self.frame, 8, 1, self.ui_state, "secrets.cloud.password") + + components.label(self.frame, 9, 0, "Cloud id", + tooltip="RUNPOD Cloud ID. The cloud service must have a public IP and SSH service. Leave empty if you want to automatically create a new RUNPOD cloud, or if you're connecting to another cloud provider via SSH Hostname and Port.") + components.entry(self.frame, 9, 1, self.ui_state, "secrets.cloud.id") + + components.label(self.frame, 10, 0, "Tensorboard TCP tunnel", + tooltip="Instead of starting tensorboard locally, make a TCP tunnel to a tensorboard on the cloud") + components.switch(self.frame, 10, 1, self.ui_state, "cloud.tensorboard_tunnel") + + + + components.label(self.frame, 1, 2, "Remote Directory", + tooltip="The directory on the cloud where files will be uploaded and downloaded.") + components.entry(self.frame, 1, 3, self.ui_state, "cloud.remote_dir") + components.label(self.frame, 2, 2, "OneTrainer Directory", + tooltip="The directory for OneTrainer on the cloud.") + components.entry(self.frame, 2, 3, self.ui_state, "cloud.onetrainer_dir") + components.label(self.frame, 3, 2, "Huggingface cache Directory", + tooltip="Huggingface models are downloaded to this remote directory.") + components.entry(self.frame, 3, 3, self.ui_state, "cloud.huggingface_cache_dir") + components.label(self.frame, 4, 2, "Install OneTrainer", + tooltip="Automatically install OneTrainer from GitHub if the directory doesn't already exist.") + components.switch(self.frame, 4, 3, self.ui_state, "cloud.install_onetrainer") + components.label(self.frame, 5, 2, "Install command", + tooltip="The command for installing OneTrainer. Leave the default, unless you want to use a development branch of OneTrainer.") + components.entry(self.frame, 5, 3, self.ui_state, "cloud.install_cmd") + components.label(self.frame, 6, 2, "Update OneTrainer", + tooltip="Update OneTrainer if it already exists on the cloud.") + components.switch(self.frame, 6, 3, self.ui_state, "cloud.update_onetrainer") + + components.label(self.frame, 8, 2, "Detach remote trainer", + tooltip="Allows the trainer to keep running even if your connection to the cloud is lost.") + components.switch(self.frame, 8, 3, self.ui_state, "cloud.detach_trainer") + components.label(self.frame, 9, 2, "Reattach id", + tooltip="An id identifying the remotely running trainer. In case you have lost connection or closed OneTrainer, it will try to reattach to this id instead of starting a new remote trainer.") + reattach_frame = ctk.CTkFrame(self.frame, fg_color="transparent") + reattach_frame.grid(row=9, column=3, padx=0, pady=0, sticky="new") + reattach_frame.grid_columnconfigure(0, weight=1) + reattach_frame.grid_columnconfigure(1, weight=1) + components.entry(reattach_frame, 0, 0, self.ui_state, "cloud.run_id", width=60) + components.button(reattach_frame, 0, 1, "Reattach now", self.__reattach) + + components.label(self.frame, 11, 2, "Download samples", + tooltip="Download samples from the remote workspace directory to your local machine.") + components.switch(self.frame, 11, 3, self.ui_state, "cloud.download_samples") + components.label(self.frame, 12, 2, "Download output model", + tooltip="Download the final model after training. You can disable this if you plan to use an automatically saved checkpoint instead.") + components.switch(self.frame, 12, 3, self.ui_state, "cloud.download_output_model") + components.label(self.frame, 13, 2, "Download saved checkpoints", + tooltip="Download the automatically saved training checkpoints from the remote workspace directory to your local machine.") + components.switch(self.frame, 13, 3, self.ui_state, "cloud.download_saves") + components.label(self.frame, 14, 2, "Download backups", + tooltip="Download backups from the remote workspace directory to your local machine. It's usually not necessary to download them, because as long as the backups are still available on the cloud, the training can be restarted using one of the cloud's backups.") + components.switch(self.frame, 14, 3, self.ui_state, "cloud.download_backups") + components.label(self.frame, 15, 2, "Download tensorboard logs", + tooltip="Download TensorBoard event logs from the remote workspace directory to your local machine. They can then be viewed locally in TensorBoard. It is recommended to disable \"Sample to TensorBoard\" to reduce the event log size.") + components.switch(self.frame, 15, 3, self.ui_state, "cloud.download_tensorboard") + components.label(self.frame, 16, 2, "Delete remote workspace", + tooltip="Delete the workspace directory on the cloud after training has finished successfully and data has been downloaded.") + components.switch(self.frame, 16, 3, self.ui_state, "cloud.delete_workspace") + + components.label(self.frame, 1, 4, "Create cloud via API", + tooltip="Automatically creates a new cloud instance if both Host:Port and Cloud ID are empty. Currently supported for RUNPOD.") + create_frame = ctk.CTkFrame(self.frame, fg_color="transparent") + create_frame.grid(row=1, column=5, padx=0, pady=0, sticky="new") + create_frame.grid_columnconfigure(0, weight=0) + create_frame.grid_columnconfigure(1, weight=1) + components.switch(create_frame, 0, 0, self.ui_state, "cloud.create") + components.button(create_frame, 0, 1, "Create cloud via website", self.__create_cloud) + + components.label(self.frame, 2, 4, "Cloud name", + tooltip="The name of the new cloud instance.") + components.entry(self.frame, 2, 5, self.ui_state, "cloud.name") + components.label(self.frame, 3, 4, "Type", + tooltip="Select the RunPod cloud type. See RunPod's website for details.") + components.options_kv(self.frame, 3, 5, [ + ("", ""), + ("Community", "COMMUNITY"), + ("Secure", "SECURE"), + ], self.ui_state, "cloud.sub_type") + + + components.label(self.frame, 4, 4, "GPU", + tooltip="Select the GPU type. Enter an API key before pressing the button.") + + _,gpu_components=components.options_adv(self.frame, 4, 5, [("")], self.ui_state, "cloud.gpu_type",adv_command=self.__set_gpu_types) + self.gpu_types_menu=gpu_components['component'] + + components.label(self.frame, 5, 4, "Volume size", + tooltip="Set the storage volume size in GB. This volume persists only until the cloud is deleted - not a RunPod network volume") + components.entry(self.frame, 5, 5, self.ui_state, "cloud.volume_size") + + components.label(self.frame, 6, 4, "Min download", + tooltip="Set the minimum download speed of the cloud in Mbps.") + components.entry(self.frame, 6, 5, self.ui_state, "cloud.min_download") + + components.label(self.frame, 8, 4, "Action on finish", + tooltip="What to do when training finishes and the data has been fully downloaded: Stop or delete the cloud, or do nothing.") + components.options_kv(self.frame, 8, 5, [ + ("None", CloudAction.NONE), + ("Stop", CloudAction.STOP), + ("Delete", CloudAction.DELETE), + ], self.ui_state, "cloud.on_finish") + + components.label(self.frame, 9, 4, "Action on error", + tooltip="What to do if training stops due to an error: Stop or delete the cloud, or do nothing. Data may be lost.") + components.options_kv(self.frame, 9, 5, [ + ("None", CloudAction.NONE), + ("Stop", CloudAction.STOP), + ("Delete", CloudAction.DELETE), + ], self.ui_state, "cloud.on_error") + + components.label(self.frame, 10, 4, "Action on detached finish", + tooltip="What to do when training finishes, but the client has been detached and cannot download data. Data may be lost.") + components.options_kv(self.frame, 10, 5, [ + ("None", CloudAction.NONE), + ("Stop", CloudAction.STOP), + ("Delete", CloudAction.DELETE), + ], self.ui_state, "cloud.on_detached_finish") + + components.label(self.frame, 11, 4, "Action on detached error", + tooltip="What to if training stops due to an error, but the client has been detached and cannot download data. Data may be lost.") + components.options_kv(self.frame, 11, 5, [ + ("None", CloudAction.NONE), + ("Stop", CloudAction.STOP), + ("Delete", CloudAction.DELETE), + ], self.ui_state, "cloud.on_detached_error") + + self.frame.pack(fill="both", expand=1) + + def __set_gpu_types(self): + self.gpu_types_menu.configure(values=[]) + if self.train_config.cloud.type == CloudType.RUNPOD: + import runpod + runpod.api_key=self.train_config.secrets.cloud.api_key + gpus=runpod.get_gpus() + self.gpu_types_menu.configure(values=[gpu['id'] for gpu in gpus]) + + def __reattach(self): + self.reattach=True + try: + self.parent.start_training() + finally: + self.reattach=False + + def __create_cloud(self): + if self.train_config.cloud.type == CloudType.RUNPOD: + webbrowser.open("https://www.runpod.io/console/deploy?template=1a33vbssq9&type=gpu", new=0, autoraise=False) diff --git a/modules/ui/CtkConceptTabView.py b/modules/ui/CtkConceptTabView.py new file mode 100644 index 000000000..0b6505694 --- /dev/null +++ b/modules/ui/CtkConceptTabView.py @@ -0,0 +1,286 @@ +import os +import pathlib +from tkinter import BooleanVar, StringVar + +from modules.ui.ConceptWindow import ConceptWindow +from modules.ui.ConfigList import ConfigList +from modules.util import path_util +from modules.util.config.ConceptConfig import ConceptConfig +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ConceptType import ConceptType +from modules.util.image_util import load_image +from modules.util.ui import components +from modules.util.ui.UIState import UIState +from modules.util.ui.validation import DebounceTimer + +import customtkinter as ctk +from PIL import Image + + +class ConceptTab(ConfigList): + + def __init__(self, master, train_config: TrainConfig, ui_state: UIState): + self.search_var = StringVar() + self.filter_var = StringVar(value="ALL") + self.show_disabled_var = BooleanVar(value=True) + + super().__init__( + master, + train_config, + ui_state, + from_external_file=True, + attr_name="concept_file_name", + config_dir="training_concepts", + default_config_name="concepts.json", + add_button_text="Add Concept", + add_button_tooltip="Adds a new concept to the current config.", + is_full_width=False, + show_toggle_button=True + ) + self._toolbar = None + self._toolbar_is_wrapped = False + self._add_search_bar() + # wrap toolbar if too narrow + self.top_frame.bind('', lambda e: self._maybe_reposition_toolbar(e.width)) + + def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): + return ConceptWidget(master, element, i, open_command, remove_command, clone_command, save_command) + + def create_new_element(self) -> dict: + return ConceptConfig.default_values() + + def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: + return ConceptWindow(self.master, self.train_config, self.current_config[i], ui_state[0], ui_state[1], ui_state[2]) + + def _add_search_bar(self): + toolbar = ctk.CTkFrame(self.top_frame, fg_color="transparent") + toolbar.grid(row=0, column=4, columnspan=2, padx=10, sticky="ew") + toolbar.grid_columnconfigure(2, weight=1) + self._toolbar = toolbar + + # Search + ctk.CTkLabel(toolbar, text="Search:").grid(row=0, column=0, padx=(0,5)) + self.search_var = StringVar() + self.search_entry = ctk.CTkEntry(toolbar, textvariable=self.search_var, + placeholder_text="Filter...", width=200) + self.search_entry.grid(row=0, column=1) + self._search_debouncer = DebounceTimer(self.search_entry, 300, lambda: self._update_filters()) + self.search_var.trace_add("write", lambda *_: self._search_debouncer.call()) + + # Spacer + ctk.CTkLabel(toolbar, text="").grid(row=0, column=2, padx=5) + + # Type filter + ctk.CTkLabel(toolbar, text="Type:").grid(row=0, column=3, padx=(0,5)) + self.filter_var = StringVar(value="ALL") + ctk.CTkOptionMenu(toolbar, values=["ALL", "STANDARD", "VALIDATION", "PRIOR_PREDICTION"], + variable=self.filter_var, command=lambda x: self._update_filters(), + width=150).grid(row=0, column=4) + + # Show disabled checkbox + self.show_disabled_var = BooleanVar(value=True) + self.show_disabled_checkbox = ctk.CTkCheckBox(toolbar, text="Show Disabled", variable=self.show_disabled_var, + command=self._update_filters, width=100) + self.show_disabled_checkbox.grid(row=0, column=5, padx=(10,0)) + self._refresh_show_disabled_text() + + # Clear button + ctk.CTkButton(toolbar, text="Clear", width=50, + command=self._reset_filters).grid(row=0, column=6, padx=(10,0)) + + def _update_filters(self): + self._create_element_list(search=self.search_var.get(), + type=self.filter_var.get(), + show_disabled=self.show_disabled_var.get()) + self._refresh_show_disabled_text() + + def _reset_filters(self): + self.search_var.set("") + self.filter_var.set("ALL") + self.show_disabled_var.set(True) + self._update_filters() + + def _element_matches_filters(self, element): + # Check enabled status + if not self.filters.get("show_disabled", True): + if hasattr(element, 'enabled') and not element.enabled: + return False + + # Search filter + search = self.filters.get("search", "").lower() + if search: + if not hasattr(element, '_search_cache'): + cache = [] + try: + if getattr(element, 'name', None): + cache.append(element.name.lower()) + p = getattr(element, 'path', None) + if p: + try: + cache.append(os.path.basename(p).lower()) + cache.append(p.lower()) + except (TypeError, AttributeError): + pass + except (AttributeError, TypeError): + pass + element._search_cache = cache + if not any(search in text for text in getattr(element, '_search_cache', [])): + return False + + # Type filter + type_filter = self.filters.get("type", "ALL") + if type_filter != "ALL": + if hasattr(element, 'type') and element.type: + try: + return ConceptType(element.type).value == type_filter + except (ValueError, AttributeError): + return False + return False + + return True + + def _maybe_reposition_toolbar(self, width): + if not self._toolbar: + return + threshold = 1070 + want_wrapped = width < threshold + if want_wrapped == self._toolbar_is_wrapped: + return + self._toolbar_is_wrapped = want_wrapped + if want_wrapped: + self._toolbar.grid_configure(row=1, column=0, columnspan=8, sticky="ew", padx=10) + else: + self._toolbar.grid_configure(row=0, column=4, columnspan=2, sticky="ew", padx=10) + + def _refresh_show_disabled_text(self): + try: + disabled_count = sum(1 for c in getattr(self, 'current_config', []) if getattr(c, 'enabled', True) is False) + except (AttributeError, TypeError): + disabled_count = 0 + text = f"Show Disabled ({disabled_count})" if disabled_count > 0 else "Show Disabled" + try: + if getattr(self, 'show_disabled_checkbox', None): + self.show_disabled_checkbox.configure(text=text) + except (AttributeError, RuntimeError): + pass + + +class ConceptWidget(ctk.CTkFrame): + def __init__(self, master, concept, i, open_command, remove_command, clone_command, save_command): + super().__init__( + master=master, width=150, height=170, corner_radius=10, bg_color="transparent" + ) + + self.concept = concept + self.ui_state = UIState(self, concept) + self.image_ui_state = UIState(self, concept.image) + self.text_ui_state = UIState(self, concept.text) + self.i = i + + self.grid_rowconfigure(1, weight=1) + + # image + self.image = ctk.CTkImage( + light_image=self.__get_preview_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) + + # name + self.name_label = components.label(self, 1, 0, self.__get_display_name(), pad=5, wraplength=140) + + # close button + close_button = ctk.CTkButton( + master=self, + width=20, + height=20, + text="X", + corner_radius=2, + fg_color="#C00000", + command=lambda: remove_command(self.i), + ) + close_button.place(x=0, y=0) + + # clone button + clone_button = ctk.CTkButton( + master=self, + width=20, + height=20, + text="+", + corner_radius=2, + fg_color="#00C000", + command=lambda: clone_command(self.i, self.__randomize_seed), + ) + clone_button.place(x=25, y=0) + + # enabled switch + enabled_switch = ctk.CTkSwitch( + master=self, + width=40, + variable=self.ui_state.get_var("enabled"), + text="", + command=save_command, + ) + enabled_switch.place(x=110, y=0) + + image_label.bind( + "", + lambda event: open_command(self.i, (self.ui_state, self.image_ui_state, self.text_ui_state)) + ) + + def __randomize_seed(self, concept: ConceptConfig): + concept.seed = ConceptConfig.default_values().seed + return concept + + def __get_display_name(self): + if self.concept.name: + return self.concept.name + elif self.concept.path: + return os.path.basename(self.concept.path) + else: + return "" + + def configure_element(self): + self.name_label.configure(text=self.__get_display_name()) + self.image.configure(light_image=self.__get_preview_image()) + try: + if hasattr(self.concept, '_search_cache'): + delattr(self.concept, '_search_cache') + except AttributeError: + pass + + def __get_preview_image(self): + preview_path = "resources/icons/icon.png" + glob_pattern = "**/*.*" if getattr(self.concept, 'include_subdirectories', False) else "*.*" + + concept_path = ConceptWindow.get_concept_path(getattr(self.concept, 'path', None)) + if concept_path: + for path in pathlib.Path(concept_path).glob(glob_pattern): + if any(part.startswith('.') for part in path.relative_to(concept_path).parent.parts): + continue + extension = os.path.splitext(path)[1] + if (path.is_file() + and path_util.is_supported_image_extension(extension) + and not path.name.endswith("-masklabel.png") + and not path.name.endswith("-condlabel.png")): + preview_path = path_util.canonical_join(concept_path, path) + break + try: + image = load_image(preview_path, convert_mode="RGBA") + except (OSError): + image = Image.new("RGBA", (150, 150), (200, 200, 200, 255)) + size = min(image.width, image.height) + image = image.crop(( + (image.width - size) // 2, + (image.height - size) // 2, + (image.width - size) // 2 + size, + (image.height - size) // 2 + size, + )) + return image.resize((150, 150), Image.Resampling.BILINEAR) + + def place_in_list(self): + index = getattr(self, 'visible_index', self.i) + x = index % 6 + y = index // 6 + self.grid(row=y, column=x, pady=5, padx=5) diff --git a/modules/ui/CtkConceptWindowView.py b/modules/ui/CtkConceptWindowView.py new file mode 100644 index 000000000..f58879d5f --- /dev/null +++ b/modules/ui/CtkConceptWindowView.py @@ -0,0 +1,934 @@ +import fractions +import math +import os +import pathlib +import platform +import random +import threading +import time +import traceback + +from modules.util import concept_stats, path_util +from modules.util.config.ConceptConfig import ConceptConfig +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.BalancingStrategy import BalancingStrategy +from modules.util.enum.ConceptType import ConceptType +from modules.util.image_util import load_image +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +from mgds.LoadingPipeline import LoadingPipeline +from mgds.OutputPipelineModule import OutputPipelineModule +from mgds.PipelineModule import PipelineModule +from mgds.pipelineModules.CapitalizeTags import CapitalizeTags +from mgds.pipelineModules.DropTags import DropTags +from mgds.pipelineModules.RandomBrightness import RandomBrightness +from mgds.pipelineModules.RandomCircularMaskShrink import ( + RandomCircularMaskShrink, +) +from mgds.pipelineModules.RandomContrast import RandomContrast +from mgds.pipelineModules.RandomFlip import RandomFlip +from mgds.pipelineModules.RandomHue import RandomHue +from mgds.pipelineModules.RandomMaskRotateCrop import RandomMaskRotateCrop +from mgds.pipelineModules.RandomRotate import RandomRotate +from mgds.pipelineModules.RandomSaturation import RandomSaturation +from mgds.pipelineModules.ShuffleTags import ShuffleTags +from mgds.pipelineModuleTypes.RandomAccessPipelineModule import ( + RandomAccessPipelineModule, +) + +import torch +from torchvision.transforms import functional + +import customtkinter as ctk +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 + + +class InputPipelineModule( + PipelineModule, + RandomAccessPipelineModule, +): + def __init__(self, data: dict): + super().__init__() + self.data = data + + def length(self) -> int: + return 1 + + def get_inputs(self) -> list[str]: + return [] + + def get_outputs(self) -> list[str]: + return list(self.data.keys()) + + def get_item(self, variation: int, index: int, requested_name: str = None) -> dict: + return self.data + + +class ConceptWindow(ctk.CTkToplevel): + def __init__( + self, + parent, + train_config: TrainConfig, + concept: ConceptConfig, + ui_state: UIState, + image_ui_state: UIState, + text_ui_state: UIState, + *args, **kwargs, + ): + super().__init__(parent, *args, **kwargs) + + self.train_config = train_config + + self.concept = concept + self.ui_state = ui_state + self.image_ui_state = image_ui_state + 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.title("Concept") + self.geometry("800x700") + self.resizable(True, True) + + self.grid_rowconfigure(0, weight=1) + self.grid_columnconfigure(0, weight=1) + + tabview = ctk.CTkTabview(self) + tabview.grid(row=0, column=0, sticky="nsew") + + self.general_tab = self.__general_tab(tabview.add("general"), concept) + self.image_augmentation_tab = self.__image_augmentation_tab(tabview.add("image augmentation")) + self.text_augmentation_tab = self.__text_augmentation_tab(tabview.add("text augmentation")) + self.concept_stats_tab = self.__concept_stats_tab(tabview.add("statistics")) + + #automatic concept scan + self.scan_thread = threading.Thread(target=self.__auto_update_concept_stats, daemon=True) + self.scan_thread.start() + + components.button(self, 1, 0, "ok", self.__ok) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def __general_tab(self, master, concept: ConceptConfig): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, weight=1) + + # name + components.label(frame, 0, 0, "Name", + tooltip="Name of the concept") + components.entry(frame, 0, 1, self.ui_state, "name") + + # enabled + components.label(frame, 1, 0, "Enabled", + tooltip="Enable or disable this concept") + components.switch(frame, 1, 1, self.ui_state, "enabled") + + # concept type + components.label(frame, 2, 0, "Concept Type", + tooltip="STANDARD: Standard finetuning with the sample as training target\n" + "VALIDATION: Use concept for validation instead of training\n" + "PRIOR_PREDICTION: Use the sample to make a prediction using the model as it was before training. This prediction is then used as the training target " + "for the model in training. This can be used as regularisation and to preserve prior model knowledge while finetuning the model on other concepts. " + "Only implemented for LoRA.", + wide_tooltip=True) + components.options(frame, 2, 1, [str(x) for x in list(ConceptType)], self.ui_state, "type") + + # path + components.label(frame, 3, 0, "Path", + tooltip="Path where the training data is located") + components.path_entry(frame, 3, 1, self.ui_state, "path", mode="dir") + 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.path_entry(frame, 4, 2, self.text_ui_state, "prompt_path", mode="file") + + def set_prompt_path_entry_enabled(option: str): + if option == 'concept': + for child in prompt_path_entry.children.values(): + child.configure(state="normal") + else: + for child in prompt_path_entry.children.values(): + child.configure(state="disabled") + + components.options_kv(frame, 4, 1, [ + ("From text file per sample", 'sample'), + ("From single text file", 'concept'), + ("From image file name", 'filename'), + ], self.text_ui_state, "prompt_source", command=set_prompt_path_entry_enabled) + set_prompt_path_entry_enabled(concept.text.prompt_source) + + # include subdirectories + components.label(frame, 5, 0, "Include Subdirectories", + tooltip="Includes images from subdirectories into the dataset") + components.switch(frame, 5, 1, self.ui_state, "include_subdirectories") + + # image variations + components.label(frame, 6, 0, "Image Variations", + tooltip="The number of different image versions to cache if latent caching is enabled.") + components.entry(frame, 6, 1, self.ui_state, "image_variations") + + # text variations + components.label(frame, 7, 0, "Text Variations", + tooltip="The number of different text versions to cache if latent caching is enabled.") + components.entry(frame, 7, 1, self.ui_state, "text_variations") + + # balancing + components.label(frame, 8, 0, "Balancing", + tooltip="The number of samples used during training. Use repeats to multiply the concept, or samples to specify an exact number of samples used in each epoch.") + components.entry(frame, 8, 1, self.ui_state, "balancing") + components.options(frame, 8, 2, [str(x) for x in list(BalancingStrategy)], self.ui_state, "balancing_strategy") + + # loss weight + components.label(frame, 9, 0, "Loss Weight", + tooltip="The loss multiplyer for this concept.") + components.entry(frame, 9, 1, self.ui_state, "loss_weight") + + frame.pack(fill="both", expand=1) + return frame + + def __image_augmentation_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=0) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # header + components.label(frame, 0, 1, "Random", + tooltip="Enable this augmentation with random values") + components.label(frame, 0, 2, "Fixed", + tooltip="Enable this augmentation with fixed values") + + # crop jitter + components.label(frame, 1, 0, "Crop Jitter", + tooltip="Enables random cropping of samples") + components.switch(frame, 1, 1, self.image_ui_state, "enable_crop_jitter") + + # random flip + components.label(frame, 2, 0, "Random Flip", + tooltip="Randomly flip the sample during training") + components.switch(frame, 2, 1, self.image_ui_state, "enable_random_flip") + components.switch(frame, 2, 2, self.image_ui_state, "enable_fixed_flip") + + # random rotation + components.label(frame, 3, 0, "Random Rotation", + tooltip="Randomly rotates the sample during training") + components.switch(frame, 3, 1, self.image_ui_state, "enable_random_rotate") + components.switch(frame, 3, 2, self.image_ui_state, "enable_fixed_rotate") + components.entry(frame, 3, 3, self.image_ui_state, "random_rotate_max_angle") + + # random brightness + components.label(frame, 4, 0, "Random Brightness", + tooltip="Randomly adjusts the brightness of the sample during training") + components.switch(frame, 4, 1, self.image_ui_state, "enable_random_brightness") + components.switch(frame, 4, 2, self.image_ui_state, "enable_fixed_brightness") + components.entry(frame, 4, 3, self.image_ui_state, "random_brightness_max_strength") + + # random contrast + components.label(frame, 5, 0, "Random Contrast", + tooltip="Randomly adjusts the contrast of the sample during training") + components.switch(frame, 5, 1, self.image_ui_state, "enable_random_contrast") + components.switch(frame, 5, 2, self.image_ui_state, "enable_fixed_contrast") + components.entry(frame, 5, 3, self.image_ui_state, "random_contrast_max_strength") + + # random saturation + components.label(frame, 6, 0, "Random Saturation", + tooltip="Randomly adjusts the saturation of the sample during training") + components.switch(frame, 6, 1, self.image_ui_state, "enable_random_saturation") + components.switch(frame, 6, 2, self.image_ui_state, "enable_fixed_saturation") + components.entry(frame, 6, 3, self.image_ui_state, "random_saturation_max_strength") + + # random hue + components.label(frame, 7, 0, "Random Hue", + tooltip="Randomly adjusts the hue of the sample during training") + components.switch(frame, 7, 1, self.image_ui_state, "enable_random_hue") + components.switch(frame, 7, 2, self.image_ui_state, "enable_fixed_hue") + components.entry(frame, 7, 3, self.image_ui_state, "random_hue_max_strength") + + # random circular mask shrink + components.label(frame, 8, 0, "Circular Mask Generation", + tooltip="Automatically create circular masks for masked training") + components.switch(frame, 8, 1, self.image_ui_state, "enable_random_circular_mask_shrink") + + # random rotate and crop + components.label(frame, 9, 0, "Random Rotate and Crop", + tooltip="Randomly rotate the training samples and crop to the masked region") + components.switch(frame, 9, 1, self.image_ui_state, "enable_random_mask_rotate_crop") + + # circular mask generation + components.label(frame, 10, 0, "Resolution Override", + tooltip="Override the resolution for this concept. Optionally specify multiple resolutions separated by a comma, or a single exact resolution in the format x") + components.switch(frame, 10, 2, self.image_ui_state, "enable_resolution_override") + components.entry(frame, 10, 3, self.image_ui_state, "resolution_override") + + # image + image_preview, filename_preview, caption_preview = self.__get_preview_image() + self.image = ctk.CTkImage( + light_image=image_preview, + size=image_preview.size, + ) + image_label = ctk.CTkLabel(master=frame, text="", image=self.image, height=300, width=300) + image_label.grid(row=0, column=4, rowspan=6) + + # refresh preview + update_button_frame = ctk.CTkFrame(master=frame, corner_radius=0, fg_color="transparent") + update_button_frame.grid(row=6, column=4, rowspan=6, sticky="nsew") + update_button_frame.grid_columnconfigure(1, weight=1) + + prev_preview_button = components.button(update_button_frame, 0, 0, "<", command=self.__prev_image_preview) + components.button(update_button_frame, 0, 1, "Update Preview", command=self.__update_image_preview) + next_preview_button = components.button(update_button_frame, 0, 2, ">", command=self.__next_image_preview) + preview_augmentations_switch = ctk.CTkSwitch(update_button_frame, text="Show Augmentations", variable=self.preview_augmentations, command=self.__update_image_preview) + preview_augmentations_switch.grid(row=1, column=0, columnspan=3, padx=5, pady=5) + + prev_preview_button.configure(width=40) + next_preview_button.configure(width=40) + + #caption and filename preview + self.filename_preview = ctk.CTkLabel(master=update_button_frame, text=filename_preview, width=300, anchor="nw", justify="left", padx=10, wraplength=280) + self.filename_preview.grid(row=2, column=0, columnspan=3) + self.caption_preview = ctk.CTkTextbox(master=update_button_frame, width = 300, height = 150, wrap="word", border_width=2) + self.caption_preview.insert(index="1.0", text=caption_preview) + self.caption_preview.configure(state="disabled") + self.caption_preview.grid(row=3, column=0, columnspan=3, rowspan=3) + + frame.pack(fill="both", expand=1) + return frame + + def __text_augmentation_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=0) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # tag shuffling + components.label(frame, 0, 0, "Tag Shuffling", + tooltip="Enables tag shuffling") + components.switch(frame, 0, 1, self.text_ui_state, "enable_tag_shuffling") + + # keep tag count + components.label(frame, 1, 0, "Tag Delimiter", + tooltip="The delimiter between tags") + components.entry(frame, 1, 1, self.text_ui_state, "tag_delimiter") + + # keep tag count + components.label(frame, 2, 0, "Keep Tag Count", + tooltip="The number of tags at the start of the caption that are not shuffled or dropped") + components.entry(frame, 2, 1, self.text_ui_state, "keep_tags_count") + + # tag dropout + components.label(frame, 3, 0, "Tag Dropout", + tooltip="Enables random dropout for tags in the captions.") + components.switch(frame, 3, 1, self.text_ui_state, "tag_dropout_enable") + components.label(frame, 4, 0, "Dropout Mode", + tooltip="Method used to drop captions. 'Full' will drop the entire caption past the 'kept' tags with a certain probability, 'Random' will drop individual tags with the set probability, and 'Random Weighted' will linearly increase the probability of dropping tags, more likely to preseve tags near the front with full probability to drop at the end.") + components.options_kv(frame, 4, 1, [ + ("Full", 'FULL'), + ("Random", 'RANDOM'), + ("Random Weighted", 'RANDOM WEIGHTED'), + ], self.text_ui_state, "tag_dropout_mode", None) + components.label(frame, 4, 2, "Probability", + tooltip="Probability to drop tags, from 0 to 1.") + components.entry(frame, 4, 3, self.text_ui_state, "tag_dropout_probability") + + components.label(frame, 5, 0, "Special Dropout Tags", + tooltip="List of tags which will be whitelisted/blacklisted by dropout. 'Whitelist' tags will never be dropped but all others may be, 'Blacklist' tags may be dropped but all others will never be, 'None' may drop any tags. Can specify either a delimiter-separated list in the field, or a file path to a .txt or .csv file with entries separated by newlines.") + components.options_kv(frame, 5, 1, [ + ("None", 'NONE'), + ("Blacklist", 'BLACKLIST'), + ("Whitelist", 'WHITELIST'), + ], self.text_ui_state, "tag_dropout_special_tags_mode", None) + components.entry(frame, 5, 2, self.text_ui_state, "tag_dropout_special_tags") + components.label(frame, 6, 0, "Special Tags Regex", + tooltip="Interpret special tags with regex, such as 'photo.*' to match 'photo, photograph, photon' but not 'telephoto'. Includes exception for '/(' and '/)' syntax found in many booru/e6 tags.") + components.switch(frame, 6, 1, self.text_ui_state, "tag_dropout_special_tags_regex") + + #capitalization randomization + components.label(frame, 7, 0, "Randomize Capitalization", + tooltip="Enables randomization of capitalization for tags in the caption.") + components.switch(frame, 7, 1, self.text_ui_state, "caps_randomize_enable") + components.label(frame, 7, 2, "Force Lowercase", + tooltip="If enabled, converts the caption to lowercase before any further processing.") + components.switch(frame, 7, 3, self.text_ui_state, "caps_randomize_lowercase") + + components.label(frame, 8, 0, "Captialization Mode", + tooltip="Comma-separated list of types of capitalization randomization to perform. 'capslock' for ALL CAPS, 'title' for First Letter Of Every Word, 'first' for First word only, 'random' for rAndOMiZeD lEtTERs.") + components.entry(frame, 8, 1, self.text_ui_state, "caps_randomize_mode") + components.label(frame, 8, 2, "Probability", + tooltip="Probability to randomize capitialization of each tag, from 0 to 1.") + components.entry(frame, 8, 3, self.text_ui_state, "caps_randomize_probability") + + frame.pack(fill="both", expand=1) + return frame + + def __concept_stats_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0, minsize=150) + frame.grid_columnconfigure(1, weight=0, minsize=150) + frame.grid_columnconfigure(2, weight=0, minsize=150) + frame.grid_columnconfigure(3, weight=0, minsize=150) + + self.cancel_scan_flag = threading.Event() + + #file size + self.file_size_label = components.label(frame, 1, 0, "Total Size", pad=0, + tooltip="Total size of all image, mask, and caption files in MB") + self.file_size_label.configure(font=ctk.CTkFont(underline=True)) + self.file_size_preview = components.label(frame, 2, 0, pad=0, text="-") + + #subdirectory count + self.dir_count_label = components.label(frame, 1, 1, "Directories", pad=0, + tooltip="Total number of directories including and under (if 'include subdirectories' is enabled) the main concept directory") + self.dir_count_label.configure(font=ctk.CTkFont(underline=True)) + self.dir_count_preview = components.label(frame, 2, 1, pad=0, text="-") + + #basic img/vid stats - count of each type in the concept + #the \n at the start of the label gives it better vertical spacing with other rows + self.image_count_label = components.label(frame, 3, 0, "\nTotal Images", pad=0, + tooltip="Total number of image files, any of the extensions " + str(path_util.SUPPORTED_IMAGE_EXTENSIONS) + ", excluding '-masklabel.png and -condlabel.png'") + self.image_count_label.configure(font=ctk.CTkFont(underline=True)) + self.image_count_preview = components.label(frame, 4, 0, pad=0, text="-") + self.video_count_label = components.label(frame, 3, 1, "\nTotal Videos", pad=0, + tooltip="Total number of video files, any of the extensions " + str(path_util.SUPPORTED_VIDEO_EXTENSIONS)) + self.video_count_label.configure(font=ctk.CTkFont(underline=True)) + self.video_count_preview = components.label(frame, 4, 1, pad=0, text="-") + self.mask_count_label = components.label(frame, 3, 2, "\nTotal Masks", pad=0, + tooltip="Total number of mask files, any file ending in '-masklabel.png'") + self.mask_count_label.configure(font=ctk.CTkFont(underline=True)) + self.mask_count_preview = components.label(frame, 4, 2, pad=0, text="-") + self.caption_count_label = components.label(frame, 3, 3, "\nTotal Captions", pad=0, + tooltip="Total number of caption files, any .txt file. With advanced scan, includes the total number of captions on separate lines across all files in parentheses.") + self.caption_count_label.configure(font=ctk.CTkFont(underline=True)) + self.caption_count_preview = components.label(frame, 4, 3, pad=0, text="-") + + #advanced img/vid stats - how many img/vid files have a mask or caption of the same name + self.image_count_mask_label = components.label(frame, 5, 0, "\nImages with Masks", pad=0, + tooltip="Total number of image files with an associated mask") + self.image_count_mask_label.configure(font=ctk.CTkFont(underline=True)) + self.image_count_mask_preview = components.label(frame, 6, 0, pad=0, text="-") + self.mask_count_label_unpaired = components.label(frame, 5, 1, "\nUnpaired Masks", pad=0, + tooltip="Total number of mask files which lack a corresponding image file - if >0, check your data set!") + self.mask_count_label_unpaired.configure(font=ctk.CTkFont(underline=True)) + self.mask_count_preview_unpaired = components.label(frame, 6, 1, pad=0, text="-") + #currently no masks for videos? + + self.image_count_caption_label = components.label(frame, 7, 0, "\nImages with Captions", pad=0, + tooltip="Total number of image files with an associated caption") + self.image_count_caption_label.configure(font=ctk.CTkFont(underline=True)) + self.image_count_caption_preview = components.label(frame, 8, 0, pad=0, text="-") + self.video_count_caption_label = components.label(frame, 7, 1, "\nVideos with Captions", pad=0, + tooltip="Total number of video files with an associated caption") + self.video_count_caption_label.configure(font=ctk.CTkFont(underline=True)) + self.video_count_caption_preview = components.label(frame, 8, 1, pad=0, text="-") + self.caption_count_label_unpaired = components.label(frame, 7, 2, "\nUnpaired Captions", pad=0, + tooltip="Total number of caption files which lack a corresponding image file - if >0, check your data set! If using 'from file name' or 'from single text file' then this can be ignored.") + self.caption_count_label_unpaired.configure(font=ctk.CTkFont(underline=True)) + self.caption_count_preview_unpaired = components.label(frame, 8, 2, pad=0, text="-") + + #resolution info + self.pixel_max_label = components.label(frame, 9, 0, "\nMax Pixels", pad=0, + tooltip="Largest image in the concept by number of pixels (width * height)") + self.pixel_max_label.configure(font=ctk.CTkFont(underline=True)) + self.pixel_max_preview = components.label(frame, 10, 0, pad=0, text="-", wraplength=150) + self.pixel_avg_label = components.label(frame, 9, 1, "\nAvg Pixels", pad=0, + tooltip="Average size of images in the concept by number of pixels (width * height)") + self.pixel_avg_label.configure(font=ctk.CTkFont(underline=True)) + self.pixel_avg_preview = components.label(frame, 10, 1, pad=0, text="-", wraplength=150) + self.pixel_min_label = components.label(frame, 9, 2, "\nMin Pixels", pad=0, + tooltip="Smallest image in the concept by number of pixels (width * height)") + self.pixel_min_label.configure(font=ctk.CTkFont(underline=True)) + self.pixel_min_preview = components.label(frame, 10, 2, pad=0, text="-", wraplength=150) + + #video length info + self.length_max_label = components.label(frame, 11, 0, "\nMax Length", pad=0, + tooltip="Longest video in the concept by number of frames") + self.length_max_label.configure(font=ctk.CTkFont(underline=True)) + self.length_max_preview = components.label(frame, 12, 0, pad=0, text="-", wraplength=150) + self.length_avg_label = components.label(frame, 11, 1, "\nAvg Length", pad=0, + tooltip="Average length of videos in the concept by number of frames") + self.length_avg_label.configure(font=ctk.CTkFont(underline=True)) + self.length_avg_preview = components.label(frame, 12, 1, pad=0, text="-", wraplength=150) + self.length_min_label = components.label(frame, 11, 2, "\nMin Length", pad=0, + tooltip="Shortest video in the concept by number of frames") + self.length_min_label.configure(font=ctk.CTkFont(underline=True)) + self.length_min_preview = components.label(frame, 12, 2, pad=0, text="-", wraplength=150) + + #video fps info + self.fps_max_label = components.label(frame, 13, 0, "\nMax FPS", pad=0, + tooltip="Video in concept with highest fps") + self.fps_max_label.configure(font=ctk.CTkFont(underline=True)) + self.fps_max_preview = components.label(frame, 14, 0, pad=0, text="-", wraplength=150) + self.fps_avg_label = components.label(frame, 13, 1, "\nAvg FPS", pad=0, + tooltip="Average fps of videos in the concept") + self.fps_avg_label.configure(font=ctk.CTkFont(underline=True)) + self.fps_avg_preview = components.label(frame, 14, 1, pad=0, text="-", wraplength=150) + self.fps_min_label = components.label(frame, 13, 2, "\nMin FPS", pad=0, + tooltip="Video in concept with the lowest fps") + self.fps_min_label.configure(font=ctk.CTkFont(underline=True)) + self.fps_min_preview = components.label(frame, 14, 2, pad=0, text="-", wraplength=150) + + #caption info + self.caption_max_label = components.label(frame, 15, 0, "\nMax Caption Length", pad=0, + tooltip="Largest caption in concept by character count. For token count, assume ~2 tokens/word") + self.caption_max_label.configure(font=ctk.CTkFont(underline=True)) + self.caption_max_preview = components.label(frame, 16, 0, pad=0, text="-", wraplength=150) + self.caption_avg_label = components.label(frame, 15, 1, "\nAvg Caption Length", pad=0, + tooltip="Average length of caption in concept by character count. For token count, assume ~2 tokens/word") + self.caption_avg_label.configure(font=ctk.CTkFont(underline=True)) + self.caption_avg_preview = components.label(frame, 16, 1, pad=0, text="-", wraplength=150) + self.caption_min_label = components.label(frame, 15, 2, "\nMin Caption Length", pad=0, + tooltip="Smallest caption in concept by character count. For token count, assume ~2 tokens/word") + self.caption_min_label.configure(font=ctk.CTkFont(underline=True)) + self.caption_min_preview = components.label(frame, 16, 2, pad=0, text="-", wraplength=150) + + #aspect bucket info + self.aspect_bucket_label = components.label(frame, 17, 0, "\nAspect Bucketing", pad=0, + tooltip="Graph of all possible buckets and the number of images in each one, defined as height/width. Buckets range from 0.25 (4:1 extremely wide) to 4 (1:4 extremely tall). \ + Images which don't match a bucket exactly are cropped to the nearest one.") + self.aspect_bucket_label.configure(font=ctk.CTkFont(underline=True)) + self.small_bucket_label = components.label(frame, 17, 1, "\nSmallest Buckets", pad=0, + tooltip="Image buckets with the least nonzero total images - if 'batch size' is larger than this, these images will be ignored during training! See the wiki for more details.") + self.small_bucket_label.configure(font=ctk.CTkFont(underline=True)) + self.small_bucket_preview = components.label(frame, 18, 1, pad=0, text="-") + + #aspect bucketing plot, mostly copied from timestep preview graph + appearance_mode = AppearanceModeTracker.get_mode() + background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) + text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) + background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" + self.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 + + assert self.bucket_fig is None + self.bucket_fig, self.bucket_ax = plt.subplots(figsize=(7,3)) + self.canvas = FigureCanvasTkAgg(self.bucket_fig, master=frame) + self.canvas.get_tk_widget().grid(row=19, column=0, columnspan=4, rowspan=2) + self.bucket_fig.tight_layout() + self.bucket_fig.subplots_adjust(bottom=0.15) + + self.bucket_fig.set_facecolor(background_color) + self.bucket_ax.set_facecolor(background_color) + self.bucket_ax.spines['bottom'].set_color(self.text_color) + self.bucket_ax.spines['left'].set_color(self.text_color) + self.bucket_ax.spines['top'].set_visible(False) + self.bucket_ax.spines['right'].set_color(self.text_color) + self.bucket_ax.tick_params(axis='x', colors=self.text_color, which="both") + self.bucket_ax.tick_params(axis='y', colors=self.text_color, which="both") + self.bucket_ax.xaxis.label.set_color(self.text_color) + self.bucket_ax.yaxis.label.set_color(self.text_color) + + #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") + self.refresh_advanced_stats_button = components.button(master=frame, row=0, column=1, text="Refresh Advanced", command=lambda: self.__get_concept_stats_threaded(True, 9999), + tooltip="Reload advanced statistics for the concept directory") #run "basic" scan first before "advanced", seems to help the system cache the directories and run faster + self.cancel_stats_button = components.button(master=frame, row=0, column=2, text="Abort Scan", command=lambda: self.__cancel_concept_stats(), + tooltip="Stop the currently running scan if it's taking a long time - advanced scan will be slow on large folders and on HDDs") + self.processing_time = components.label(frame, 0, 3, text="-", tooltip="Time taken to process concept directory") + + frame.pack(fill="both", expand=1) + return frame + + def __prev_image_preview(self): + self.image_preview_file_index = max(self.image_preview_file_index - 1, 0) + self.__update_image_preview() + + def __next_image_preview(self): + self.image_preview_file_index += 1 + self.__update_image_preview() + + def __update_image_preview(self): + image_preview, filename_preview, caption_preview = self.__get_preview_image() + self.image.configure(light_image=image_preview, size=image_preview.size) + self.filename_preview.configure(text=filename_preview) + self.caption_preview.configure(state="normal") + self.caption_preview.delete(index1="1.0", index2="end") + self.caption_preview.insert(index="1.0", text=caption_preview) + self.caption_preview.configure(state="disabled") + + @staticmethod + def get_concept_path(path: str) -> str | None: + if os.path.isdir(path): + return path + try: + #don't download, only check if available locally: + return huggingface_hub.snapshot_download(repo_id=path, repo_type="dataset", local_files_only=True) + except Exception: + return None + + def __download_dataset(self): + try: + huggingface_hub.login(token=self.train_config.secrets.huggingface_token, new_session=False) + huggingface_hub.snapshot_download(repo_id=self.concept.path, repo_type="dataset") + except Exception: + traceback.print_exc() + + def __download_dataset_threaded(self): + download_thread = threading.Thread(target=self.__download_dataset, daemon=True) + download_thread.start() + + def _read_text_file_for_preview(self, file_path: str) -> str: + empty_msg = "[Empty prompt]" + try: + with open(file_path, "r") as f: + if self.preview_augmentations.get(): + lines = [line.strip() for line in f if line.strip()] + return random.choice(lines) if lines else empty_msg + content = f.read().strip() + return content if content else empty_msg + except FileNotFoundError: + return "File not found, please check the path" + except IsADirectoryError: + return "[Provided path is a directory, please correct the caption path]" + except PermissionError: + if platform.system() == "Windows": + return "[Permission denied, please check the file permissions or Windows Defender settings]" + else: + return "[Permission denied, please check the file permissions]" + except UnicodeDecodeError: + return "[Invalid file encoding. This should not happen, please report this issue]" + + def __get_preview_image(self): + preview_image_path = "resources/icons/icon.png" + file_index = -1 + glob_pattern = "**/*.*" if self.concept.include_subdirectories else "*.*" + + concept_path = self.get_concept_path(self.concept.path) + if concept_path: + for path in pathlib.Path(concept_path).glob(glob_pattern): + if any(part.startswith('.') for part in path.relative_to(concept_path).parent.parts): + continue + extension = os.path.splitext(path)[1] + if path.is_file() and path_util.is_supported_image_extension(extension) \ + and not path.name.endswith("-masklabel.png") and not path.name.endswith("-condlabel.png"): + preview_image_path = path_util.canonical_join(concept_path, path) + file_index += 1 + if file_index == self.image_preview_file_index: + break + + image = load_image(preview_image_path, 'RGB') + image_tensor = functional.to_tensor(image) + + splitext = os.path.splitext(preview_image_path) + preview_mask_path = path_util.canonical_join(splitext[0] + "-masklabel.png") + if not os.path.isfile(preview_mask_path): + preview_mask_path = None + + if preview_mask_path: + mask = Image.open(preview_mask_path).convert("L") + mask_tensor = functional.to_tensor(mask) + else: + mask_tensor = torch.ones((1, image_tensor.shape[1], image_tensor.shape[2])) + + source = self.concept.text.prompt_source + preview_p = pathlib.Path(preview_image_path) + if source == "filename": + prompt_output = preview_p.stem or "[Empty prompt]" + else: + file_map = { + "sample": preview_p.with_suffix(".txt"), + "concept": pathlib.Path(self.concept.text.prompt_path) if self.concept.text.prompt_path else None, + } + file_path = file_map.get(source) + prompt_output = self._read_text_file_for_preview(str(file_path)) if file_path else "[Empty prompt]" + + modules = [] + if self.preview_augmentations.get(): + input_module = InputPipelineModule({ + 'true': True, + 'image': image_tensor, + 'mask': mask_tensor, + 'enable_random_flip': self.concept.image.enable_random_flip, + 'enable_fixed_flip': self.concept.image.enable_fixed_flip, + 'enable_random_rotate': self.concept.image.enable_random_rotate, + 'enable_fixed_rotate': self.concept.image.enable_fixed_rotate, + 'random_rotate_max_angle': self.concept.image.random_rotate_max_angle, + 'enable_random_brightness': self.concept.image.enable_random_brightness, + 'enable_fixed_brightness': self.concept.image.enable_fixed_brightness, + 'random_brightness_max_strength': self.concept.image.random_brightness_max_strength, + 'enable_random_contrast': self.concept.image.enable_random_contrast, + 'enable_fixed_contrast': self.concept.image.enable_fixed_contrast, + 'random_contrast_max_strength': self.concept.image.random_contrast_max_strength, + 'enable_random_saturation': self.concept.image.enable_random_saturation, + 'enable_fixed_saturation': self.concept.image.enable_fixed_saturation, + 'random_saturation_max_strength': self.concept.image.random_saturation_max_strength, + 'enable_random_hue': self.concept.image.enable_random_hue, + 'enable_fixed_hue': self.concept.image.enable_fixed_hue, + 'random_hue_max_strength': self.concept.image.random_hue_max_strength, + 'enable_random_circular_mask_shrink': self.concept.image.enable_random_circular_mask_shrink, + 'enable_random_mask_rotate_crop': self.concept.image.enable_random_mask_rotate_crop, + + 'prompt' : prompt_output, + 'tag_dropout_enable' : self.concept.text.tag_dropout_enable, + 'tag_dropout_probability' : self.concept.text.tag_dropout_probability, + 'tag_dropout_mode' : self.concept.text.tag_dropout_mode, + 'tag_dropout_special_tags' : self.concept.text.tag_dropout_special_tags, + 'tag_dropout_special_tags_mode' : self.concept.text.tag_dropout_special_tags_mode, + 'tag_delimiter' : self.concept.text.tag_delimiter, + 'keep_tags_count' : self.concept.text.keep_tags_count, + 'tag_dropout_special_tags_regex' : self.concept.text.tag_dropout_special_tags_regex, + 'caps_randomize_enable' : self.concept.text.caps_randomize_enable, + 'caps_randomize_probability' : self.concept.text.caps_randomize_probability, + 'caps_randomize_mode' : self.concept.text.caps_randomize_mode, + 'caps_randomize_lowercase' : self.concept.text.caps_randomize_lowercase, + 'enable_tag_shuffling' : self.concept.text.enable_tag_shuffling, + }) + + circular_mask_shrink = RandomCircularMaskShrink(mask_name='mask', shrink_probability=1.0, shrink_factor_min=0.2, shrink_factor_max=1.0, enabled_in_name='enable_random_circular_mask_shrink') + random_mask_rotate_crop = RandomMaskRotateCrop(mask_name='mask', additional_names=['image'], min_size=512, min_padding_percent=10, max_padding_percent=30, max_rotate_angle=20, enabled_in_name='enable_random_mask_rotate_crop') + random_flip = RandomFlip(names=['image', 'mask'], enabled_in_name='enable_random_flip', fixed_enabled_in_name='enable_fixed_flip') + random_rotate = RandomRotate(names=['image', 'mask'], enabled_in_name='enable_random_rotate', fixed_enabled_in_name='enable_fixed_rotate', max_angle_in_name='random_rotate_max_angle') + random_brightness = RandomBrightness(names=['image'], enabled_in_name='enable_random_brightness', fixed_enabled_in_name='enable_fixed_brightness', max_strength_in_name='random_brightness_max_strength') + random_contrast = RandomContrast(names=['image'], enabled_in_name='enable_random_contrast', fixed_enabled_in_name='enable_fixed_contrast', max_strength_in_name='random_contrast_max_strength') + random_saturation = RandomSaturation(names=['image'], enabled_in_name='enable_random_saturation', fixed_enabled_in_name='enable_fixed_saturation', max_strength_in_name='random_saturation_max_strength') + random_hue = RandomHue(names=['image'], enabled_in_name='enable_random_hue', fixed_enabled_in_name='enable_fixed_hue', max_strength_in_name='random_hue_max_strength') + drop_tags = DropTags(text_in_name='prompt', enabled_in_name='tag_dropout_enable', probability_in_name='tag_dropout_probability', dropout_mode_in_name='tag_dropout_mode', + special_tags_in_name='tag_dropout_special_tags', special_tag_mode_in_name='tag_dropout_special_tags_mode', delimiter_in_name='tag_delimiter', + keep_tags_count_in_name='keep_tags_count', text_out_name='prompt', regex_enabled_in_name='tag_dropout_special_tags_regex') + caps_randomize = CapitalizeTags(text_in_name='prompt', enabled_in_name='caps_randomize_enable', probability_in_name='caps_randomize_probability', + capitalize_mode_in_name='caps_randomize_mode', delimiter_in_name='tag_delimiter', convert_lowercase_in_name='caps_randomize_lowercase', text_out_name='prompt') + shuffle_tags = ShuffleTags(text_in_name='prompt', enabled_in_name='enable_tag_shuffling', delimiter_in_name='tag_delimiter', keep_tags_count_in_name='keep_tags_count', text_out_name='prompt') + output_module = OutputPipelineModule(['image', 'mask', 'prompt']) + + modules = [ + input_module, + circular_mask_shrink, + random_mask_rotate_crop, + random_flip, + random_rotate, + random_brightness, + random_contrast, + random_saturation, + random_hue, + drop_tags, + caps_randomize, + shuffle_tags, + output_module, + ] + + pipeline = LoadingPipeline( + device=torch.device('cpu'), + modules=modules, + batch_size=1, + seed=random.randint(0, 2**30), + state=None, + initial_epoch=0, + initial_index=0, + ) + + data = pipeline.__next__() + image_tensor = data['image'] + mask_tensor = data['mask'] + prompt_output = data['prompt'] + + filename_output = os.path.basename(preview_image_path) + + mask_tensor = torch.clamp(mask_tensor, 0.3, 1) + image_tensor = image_tensor * mask_tensor + + image = functional.to_pil_image(image_tensor) + + image.thumbnail((300, 300)) + + return image, filename_output, prompt_output + + def __update_concept_stats(self): + #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") + + #directory count + self.dir_count_preview.configure(text=self.concept.concept_stats["directory_count"]) + + #image count + self.image_count_preview.configure(text=self.concept.concept_stats["image_count"]) + self.image_count_mask_preview.configure(text=self.concept.concept_stats["image_with_mask_count"]) + self.image_count_caption_preview.configure(text=self.concept.concept_stats["image_with_caption_count"]) + + #video count + self.video_count_preview.configure(text=self.concept.concept_stats["video_count"]) + #self.video_count_mask_preview.configure(text=self.concept.concept_stats["video_with_mask_count"]) + self.video_count_caption_preview.configure(text=self.concept.concept_stats["video_with_caption_count"]) + + #mask count + self.mask_count_preview.configure(text=self.concept.concept_stats["mask_count"]) + self.mask_count_preview_unpaired.configure(text=self.concept.concept_stats["unpaired_masks"]) + + #caption count + if self.concept.concept_stats["subcaption_count"] > 0: + self.caption_count_preview.configure(text=f'{self.concept.concept_stats["caption_count"]} ({self.concept.concept_stats["subcaption_count"]})') + else: + self.caption_count_preview.configure(text=self.concept.concept_stats["caption_count"]) + self.caption_count_preview_unpaired.configure(text=self.concept.concept_stats["unpaired_captions"]) + + #resolution info + max_pixels = self.concept.concept_stats["max_pixels"] + avg_pixels = self.concept.concept_stats["avg_pixels"] + min_pixels = self.concept.concept_stats["min_pixels"] + + if any(isinstance(x, str) for x in [max_pixels, avg_pixels, min_pixels]) or self.concept.concept_stats["image_count"] == 0: #will be str if adv stats were not taken + self.pixel_max_preview.configure(text="-") + self.pixel_avg_preview.configure(text="-") + self.pixel_min_preview.configure(text="-") + else: + #formatted as (#pixels/1000000) MP, width x height, \n filename + self.pixel_max_preview.configure(text=f'{str(round(max_pixels[0]/1000000, 2))} MP, {max_pixels[2]}\n{max_pixels[1]}') + self.pixel_avg_preview.configure(text=f'{str(round(avg_pixels/1000000, 2))} MP, ~{int(math.sqrt(avg_pixels))}w x {int(math.sqrt(avg_pixels))}h') + self.pixel_min_preview.configure(text=f'{str(round(min_pixels[0]/1000000, 2))} MP, {min_pixels[2]}\n{min_pixels[1]}') + + #video length and fps info + max_length = self.concept.concept_stats["max_length"] + avg_length = self.concept.concept_stats["avg_length"] + min_length = self.concept.concept_stats["min_length"] + max_fps = self.concept.concept_stats["max_fps"] + avg_fps = self.concept.concept_stats["avg_fps"] + min_fps = self.concept.concept_stats["min_fps"] + + if any(isinstance(x, str) for x in [max_length, avg_length, min_length]) or self.concept.concept_stats["video_count"] == 0: #will be str if adv stats were not taken + self.length_max_preview.configure(text="-") + self.length_avg_preview.configure(text="-") + self.length_min_preview.configure(text="-") + self.fps_max_preview.configure(text="-") + self.fps_avg_preview.configure(text="-") + self.fps_min_preview.configure(text="-") + else: + #formatted as (#frames) frames \n filename + self.length_max_preview.configure(text=f'{int(max_length[0])} frames\n{max_length[1]}') + self.length_avg_preview.configure(text=f'{int(avg_length)} frames') + self.length_min_preview.configure(text=f'{int(min_length[0])} frames\n{min_length[1]}') + #formatted as (#fps) fps \n filename + self.fps_max_preview.configure(text=f'{int(max_fps[0])} fps\n{max_fps[1]}') + self.fps_avg_preview.configure(text=f'{int(avg_fps)} fps') + self.fps_min_preview.configure(text=f'{int(min_fps[0])} fps\n{min_fps[1]}') + + #caption info + max_caption_length = self.concept.concept_stats["max_caption_length"] + avg_caption_length = self.concept.concept_stats["avg_caption_length"] + min_caption_length = self.concept.concept_stats["min_caption_length"] + + if any(isinstance(x, str) for x in [max_caption_length, avg_caption_length, min_caption_length]) or self.concept.concept_stats["caption_count"] == 0: #will be str if adv stats were not taken + self.caption_max_preview.configure(text="-") + self.caption_avg_preview.configure(text="-") + self.caption_min_preview.configure(text="-") + else: + #formatted as (#chars) chars, (#words) words, \n filename + self.caption_max_preview.configure(text=f'{max_caption_length[0]} chars, {max_caption_length[2]} words\n{max_caption_length[1]}') + self.caption_avg_preview.configure(text=f'{int(avg_caption_length[0])} chars, {int(avg_caption_length[1])} words') + self.caption_min_preview.configure(text=f'{min_caption_length[0]} chars, {min_caption_length[2]} words\n{min_caption_length[1]}') + + #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) + + self.bucket_ax.cla() + aspects = [str(x) for x in list(aspect_buckets.keys())] + aspect_ratios = [self.decimal_to_aspect_ratio(x) for x in list(aspect_buckets.keys())] + counts = list(aspect_buckets.values()) + b = self.bucket_ax.bar(aspect_ratios, counts) + self.bucket_ax.bar_label(b, color=self.text_color) + sec = self.bucket_ax.secondary_xaxis(location=-0.1) + sec.spines["bottom"].set_linewidth(0) + sec.set_xticks([0, (len(aspects)-1)/2, len(aspects)-1], labels=["Wide", "Square", "Tall"]) + sec.tick_params('x', length=0) + self.canvas.draw() + + def decimal_to_aspect_ratio(self, value : float): + #find closest fraction to decimal aspect value and convert to a:b format + aspect_fraction = fractions.Fraction(value).limit_denominator(16) + aspect_string = f'{aspect_fraction.denominator}:{aspect_fraction.numerator}' + return aspect_string + + def __get_concept_stats(self, advanced_checks: bool, wait_time: float): + start_time = time.perf_counter() + last_update = time.perf_counter() + self.cancel_scan_flag.clear() + self.concept_stats_tab.after(0, self.__disable_scan_buttons) + concept_path = self.get_concept_path(self.concept.path) + + if not concept_path: + print(f"Unable to get statistics for concept path: {self.concept.path}") + self.concept_stats_tab.after(0, self.__enable_scan_buttons) + return + subfolders = [concept_path] + + 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: + 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 + subfolders.extend([f for f in os.scandir(path) if f.is_dir() and not f.name.startswith('.')]) + self.concept.concept_stats = stats_dict + #update GUI approx every half second + if time.perf_counter() > (last_update + 0.5): + 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) + + 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): + self.refresh_basic_stats_button.configure(state="disabled") + self.refresh_advanced_stats_button.configure(state="disabled") + + def __enable_scan_buttons(self): + self.refresh_basic_stats_button.configure(state="normal") + self.refresh_advanced_stats_button.configure(state="normal") + + def __cancel_concept_stats(self): + self.cancel_scan_flag.set() + + def __auto_update_concept_stats(self): + try: + self.__update_concept_stats() #load stats from config if available, else raises KeyError + if self.concept.concept_stats["file_size"] == 0: #force rescan if empty + raise KeyError + except KeyError: + concept_path = self.get_concept_path(self.concept.path) + if concept_path: + self.__get_concept_stats(False, 2) #force rescan if config is empty, timeout of 2 sec + if self.concept.concept_stats["processing_time"] < 0.1: + self.__get_concept_stats(True, 2) #do advanced scan automatically if basic took <0.1s + + def destroy(self): + if self.bucket_fig is not None: + plt.close(self.bucket_fig) + self.bucket_fig = None + + super().destroy() + + def __ok(self): + self.destroy() diff --git a/modules/ui/CtkConfigListView.py b/modules/ui/CtkConfigListView.py new file mode 100644 index 000000000..75d69252a --- /dev/null +++ b/modules/ui/CtkConfigListView.py @@ -0,0 +1,354 @@ +import contextlib +import copy +import json +import os +import tkinter as tk +from abc import ABCMeta, abstractmethod + +from modules.util import path_util +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.UIState import UIState + +import customtkinter as ctk + + +class ConfigList(metaclass=ABCMeta): + + def __init__( + self, + master, + train_config: TrainConfig, + ui_state: UIState, + from_external_file: bool, + attr_name: str = "", + enable_key: str = "enabled", + config_dir: str = "", + default_config_name: str = "", + add_button_text: str = "", + add_button_tooltip: str = "", + is_full_width: bool = "", + show_toggle_button: bool = False, + ): + self.master = master + self.train_config = train_config + self.ui_state = ui_state + self.from_external_file = from_external_file + self.attr_name = attr_name + self.enable_key = enable_key + + self.config_dir = config_dir + self.default_config_name = default_config_name + + self.is_full_width = is_full_width + + # From search-concepts + self.filters = {"search": "", "type": "ALL", "show_disabled": True} + self.widgets_initialized = False + + # From master + self.toggle_button = None + self.show_toggle_button = show_toggle_button + self.is_opening_window = False + self._is_current_item_enabled = False + + self.master.grid_rowconfigure(0, weight=0) + self.master.grid_rowconfigure(1, weight=1) + self.master.grid_columnconfigure(0, weight=1) + + if self.from_external_file: + self.top_frame = ctk.CTkFrame(self.master, fg_color="transparent") + self.top_frame.grid(row=0, column=0, sticky="nsew") + + self.configs_dropdown = None + self.element_list = None + + self.configs = [] + self.__load_available_config_names() + + 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) + 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.element_list = None + 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() + + + + @abstractmethod + def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): + pass + + @abstractmethod + def create_new_element(self) -> BaseConfig: + pass + + @abstractmethod + def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: + pass + + def _refresh_show_disabled_text(self): + return + + def _reset_filters(self): # pragma: no cover - default noop + search_var = getattr(self, 'search_var', None) + filter_var = getattr(self, 'filter_var', None) + show_disabled_var = getattr(self, 'show_disabled_var', None) + + if search_var: + search_var.set("") + if filter_var: + filter_var.set("ALL") + if show_disabled_var: + show_disabled_var.set(True) + if search_var and hasattr(self, '_update_filters'): + self._update_filters() + + def _update_item_enabled_state(self): + # Only count items that match current filters + self._is_current_item_enabled = any( + item.ui_state.get_var(self.enable_key).get() + for i, item in enumerate(self.widgets) + if i < len(self.current_config) and self._element_matches_filters(self.current_config[i]) + ) + + def _update_toggle_button_text(self): + if not self.show_toggle_button: + return + self._update_item_enabled_state() + if self.toggle_button is not None: + self.toggle_button.configure(text="Disable" if self._is_current_item_enabled else "Enable") + + def _toggle(self): + self._toggle_items() + + def _toggle_items(self): + enable_state = not self._is_current_item_enabled + + # Only toggle items that match current filters + for i, widget in enumerate(self.widgets): + if i < len(self.current_config) and self._element_matches_filters(self.current_config[i]): + widget.ui_state.get_var(self.enable_key).set(enable_state) + self.save_current_config() + + self._update_widget_visibility() + + def __create_configs_dropdown(self): + if self.configs_dropdown is not None: + self.configs_dropdown.destroy() + + self.configs_dropdown = components.options_kv( + self.top_frame, 0, 0, self.configs, self.ui_state, self.attr_name, self.__load_current_config + ) + self._update_toggle_button_text() + + def _create_element_list(self, **filters): + if not self.from_external_file: + self.current_config = getattr(self.train_config, self.attr_name) + + self.filters.update(filters) + + if not self.widgets_initialized: + self._initialize_all_widgets() + self.widgets_initialized = True + + self._update_widget_visibility() + self._update_toggle_button_text() + + def _initialize_all_widgets(self): + self.widgets = [] + if self.element_list is not None: + self.element_list.destroy() + + self.element_list = ctk.CTkScrollableFrame(self.master, fg_color="transparent") + self.element_list.grid(row=1, column=0, sticky="nsew") + + if self.is_full_width: + self.element_list.grid_columnconfigure(0, weight=1) + + for i, element in enumerate(self.current_config): + widget = self.create_widget( + self.element_list, element, i, + self.__open_element_window, + self.__remove_element, + self.__clone_element, + self.save_current_config + ) + self.widgets.append(widget) + + def _update_widget_visibility(self): + visible_index = 0 + + for i, widget in enumerate(self.widgets): + if i < len(self.current_config): + element = self.current_config[i] + + if self._element_matches_filters(element): + widget.visible_index = visible_index + widget.place_in_list() + visible_index += 1 + else: + widget.grid_remove() + + def __load_available_config_names(self): + if os.path.isdir(self.config_dir): + for path in os.listdir(self.config_dir): + path = path_util.canonical_join(self.config_dir, path) + if path.endswith(".json") and os.path.isfile(path): + name = os.path.basename(path) + name = os.path.splitext(name)[0] + self.configs.append((name, path)) + + if len(self.configs) == 0: + name = self.default_config_name.removesuffix(".json") + self.__create_config(name) + self.save_current_config() + + def __create_config(self, name: str): + name = path_util.safe_filename(name) + path = path_util.canonical_join(self.config_dir, f"{name}.json") + self.configs.append((name, path)) + self.__create_configs_dropdown() + + def __add_config(self): + dialogs.StringInputDialog(self.master, "name", "Name", self.__create_config) + + def __add_element(self): + new_element = self.create_new_element() + 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 + widget = self.create_widget( + self.element_list, new_element, i, + self.__open_element_window, + self.__remove_element, + self.__clone_element, + self.save_current_config + ) + self.widgets.append(widget) + self._update_widget_visibility() + else: + self.widgets_initialized = False + self._create_element_list() + self.save_current_config() + + def __clone_element(self, clone_i, modify_element_fun=None): + new_element = copy.deepcopy(self.current_config[clone_i]) + + if modify_element_fun is not None: + new_element = modify_element_fun(new_element) + self.current_config.append(new_element) + if self.widgets_initialized and self.element_list is not None: + i = len(self.current_config) - 1 + widget = self.create_widget( + self.element_list, new_element, i, + self.__open_element_window, + self.__remove_element, + self.__clone_element, + self.save_current_config + ) + self.widgets.append(widget) + self._update_widget_visibility() + else: + self.widgets_initialized = False + self._create_element_list() + self.save_current_config() + + 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() + else: + self.widgets_initialized = False + self._create_element_list() + self.save_current_config() + + def __load_current_config(self, filename): + try: + with open(filename, "r") as f: + self.current_config = [] + + loaded_config_json = json.load(f) + for element_json in loaded_config_json: + element = self.create_new_element().from_dict(element_json) + self.current_config.append(element) + except (FileNotFoundError, json.JSONDecodeError) as e: + print(f"Failed to load config from {filename}: {e}") + self.current_config = [] + + # reset filters when switching configs + if hasattr(self, '_reset_filters') and self.widgets_initialized: + self._reset_filters() + + self.widgets_initialized = False + self._create_element_list() + self._update_toggle_button_text() + + def save_current_config(self): + if self.from_external_file: + try: + if not os.path.exists(self.config_dir): + os.makedirs(self.config_dir, exist_ok=True) + + write_json_atomic( + getattr(self.train_config, self.attr_name), + [element.to_dict() for element in self.current_config] + ) + except (OSError) as e: + print(f"Failed to save config: {e}") + + self._update_toggle_button_text() + + if self.widgets_initialized: + try: + self._update_widget_visibility() + except (tk.TclError, AttributeError) as e: + print.debug(f"Widget visibility update failed: {e}") + + # let subclass refresh any show-disabled UI + if hasattr(self, '_refresh_show_disabled_text'): + self._refresh_show_disabled_text() + + def _element_matches_filters(self, element): + return True # Show all by default + + def __open_element_window(self, i, ui_state): + if self.is_opening_window: + return + self.is_opening_window = True + try: + window = self.open_element_window(i, ui_state) + self.master.wait_window(window) + try: + if self.widgets is not None and 0 <= i < len(self.widgets): + self.widgets[i].configure_element() + except Exception: + self.widgets_initialized = False + self._create_element_list() + self.save_current_config() + finally: + self.is_opening_window = False diff --git a/modules/ui/CtkConvertModelUIView.py b/modules/ui/CtkConvertModelUIView.py new file mode 100644 index 000000000..6cb1b507a --- /dev/null +++ b/modules/ui/CtkConvertModelUIView.py @@ -0,0 +1,170 @@ +import traceback +from uuid import uuid4 + +from modules.util import create +from modules.util.args.ConvertModelArgs import ConvertModelArgs +from modules.util.config.TrainConfig import QuantizationConfig +from modules.util.enum.DataType import DataType +from modules.util.enum.ModelFormat import ModelFormat +from modules.util.enum.ModelType import ModelType +from modules.util.enum.PathIOType import PathIOType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.ModelNames import EmbeddingName, ModelNames +from modules.util.torch_util import torch_gc +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class ConvertModelUI(ctk.CTkToplevel): + def __init__(self, parent, *args, **kwargs): + super().__init__(parent, *args, **kwargs) + self.parent = parent + + self.parent = parent + self.convert_model_args = ConvertModelArgs.default_values() + self.ui_state = UIState(self, self.convert_model_args) + self.button = None + + + self.title("Convert models") + self.geometry("550x350") + self.resizable(True, True) + + self.frame = ctk.CTkFrame(self, width=600, height=300) + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + self.frame.grid_columnconfigure(0, weight=0) + self.frame.grid_columnconfigure(1, weight=1) + + self.main_frame(self.frame) + self.frame.pack(fill="both", expand=True) + + self.wait_visibility() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def main_frame(self, master): + # model type + components.label(master, 0, 0, "Model Type", + tooltip="Type of the model") + components.options_kv(master, 0, 1, [ #TODO simplify + ("Stable Diffusion 1.5", ModelType.STABLE_DIFFUSION_15), + ("Stable Diffusion 1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), + ("Stable Diffusion 2.0", ModelType.STABLE_DIFFUSION_20), + ("Stable Diffusion 2.0 Inpainting", ModelType.STABLE_DIFFUSION_20_INPAINTING), + ("Stable Diffusion 2.1", ModelType.STABLE_DIFFUSION_21), + ("Stable Diffusion 3", ModelType.STABLE_DIFFUSION_3), + ("Stable Diffusion 3.5", ModelType.STABLE_DIFFUSION_35), + ("Stable Diffusion XL 1.0 Base", ModelType.STABLE_DIFFUSION_XL_10_BASE), + ("Stable Diffusion XL 1.0 Base Inpainting", ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING), + ("Wuerstchen v2", ModelType.WUERSTCHEN_2), + ("Stable Cascade", ModelType.STABLE_CASCADE_1), + ("PixArt Alpha", ModelType.PIXART_ALPHA), + ("PixArt Sigma", ModelType.PIXART_SIGMA), + ("Flux Dev", ModelType.FLUX_DEV_1), + ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), + ("Flux 2", ModelType.FLUX_2), + ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), + ("Chroma1", ModelType.CHROMA_1), #TODO does this just work? HiDream is not here + ("QwenImage", ModelType.QWEN), #TODO does this just work? HiDream is not here + ("ZImage", ModelType.Z_IMAGE), + ], self.ui_state, "model_type") + + # training method + components.label(master, 1, 0, "Model Type", + tooltip="The type of model to convert") + components.options_kv(master, 1, 1, [ + ("Base Model", TrainingMethod.FINE_TUNE), + ("LoRA", TrainingMethod.LORA), + ("Embedding", TrainingMethod.EMBEDDING), + ], self.ui_state, "training_method") + + # input name + components.label(master, 2, 0, "Input name", + tooltip="Filename, directory or hugging face repository of the base model") + components.path_entry( + master, 2, 1, self.ui_state, "input_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # output data type + components.label(master, 3, 0, "Output Data Type", + tooltip="Precision to use when saving the output model") + components.options_kv(master, 3, 1, [ + ("float32", DataType.FLOAT_32), + ("float16", DataType.FLOAT_16), + ("bfloat16", DataType.BFLOAT_16), + ], self.ui_state, "output_dtype") + + # output format + components.label(master, 4, 0, "Output Format", + tooltip="Format to use when saving the output model") + components.options_kv(master, 4, 1, [ + ("Safetensors", ModelFormat.SAFETENSORS), + ("Diffusers", ModelFormat.DIFFUSERS), + ], self.ui_state, "output_model_format") + + # output model destination + components.label(master, 5, 0, "Model Output Destination", + tooltip="Filename or directory where the output model is saved") + components.path_entry( + master, 5, 1, self.ui_state, "output_model_destination", + mode="file", + io_type=PathIOType.MODEL, + ) + + self.button = components.button(master, 6, 1, "Convert", self.convert_model) + + def convert_model(self): + try: + self.button.configure(state="disabled") + model_loader = create.create_model_loader( + model_type=self.convert_model_args.model_type, + training_method=self.convert_model_args.training_method + ) + model_saver = create.create_model_saver( + model_type=self.convert_model_args.model_type, + training_method=self.convert_model_args.training_method + ) + + print("Loading model " + self.convert_model_args.input_name) + if self.convert_model_args.training_method in [TrainingMethod.FINE_TUNE]: + model = model_loader.load( + model_type=self.convert_model_args.model_type, + model_names=ModelNames( + base_model=self.convert_model_args.input_name, + ), + weight_dtypes=self.convert_model_args.weight_dtypes(), + quantization=QuantizationConfig.default_values(), + ) + elif self.convert_model_args.training_method in [TrainingMethod.LORA, TrainingMethod.EMBEDDING]: + model = model_loader.load( + model_type=self.convert_model_args.model_type, + model_names=ModelNames( + base_model=None, + lora=self.convert_model_args.input_name, + embedding=EmbeddingName(str(uuid4()), self.convert_model_args.input_name), + ), + weight_dtypes=self.convert_model_args.weight_dtypes(), + quantization=QuantizationConfig.default_values(), + ) + else: + raise Exception("could not load model: " + self.convert_model_args.input_name) + + print("Saving model " + self.convert_model_args.output_model_destination) + model_saver.save( + model=model, + model_type=self.convert_model_args.model_type, + output_model_format=self.convert_model_args.output_model_format, + output_model_destination=self.convert_model_args.output_model_destination, + dtype=self.convert_model_args.output_dtype.torch_dtype(), + ) + print("Model converted") + except Exception: + traceback.print_exc() + + torch_gc() + self.button.configure(state="normal") diff --git a/modules/ui/CtkGenerateCaptionsWindowView.py b/modules/ui/CtkGenerateCaptionsWindowView.py new file mode 100644 index 000000000..1690879f1 --- /dev/null +++ b/modules/ui/CtkGenerateCaptionsWindowView.py @@ -0,0 +1,133 @@ +import contextlib +import tkinter as tk +from tkinter import filedialog + +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk + + +class GenerateCaptionsWindow(ctk.CTkToplevel): + def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs): + """ + Window for generating captions for a folder of images + + Parameters: + parent (`Tk`): the parent window + path (`str`): the path to the folder + parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox + """ + super().__init__(parent, *args, **kwargs) + self.parent = parent + + if path is None: + path = "" + + self.mode_var = ctk.StringVar(self, "Create if absent") + self.modes = ["Replace all captions", "Create if absent", "Add as new line"] + self.model_var = ctk.StringVar(self, "Blip") + self.models = ["Blip", "Blip2", "WD14 VIT v2"] + + self.title("Batch generate captions") + self.geometry("360x360") + self.resizable(True, True) + + self.frame = ctk.CTkFrame(self, width=600, height=300) + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + + self.model_label = ctk.CTkLabel(self.frame, text="Model", width=100) + self.model_label.grid(row=0, column=0, sticky="w", padx=5, pady=5) + self.model_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.model_var, values=self.models, dynamic_resizing=False, width=200) + self.model_dropdown.grid(row=0, column=1, sticky="w", padx=5, pady=5) + + self.path_label = ctk.CTkLabel(self.frame, text="Folder", width=100) + self.path_label.grid(row=1, column=0, sticky="w",padx=5, pady=5) + self.path_entry = ctk.CTkEntry(self.frame, width=150) + self.path_entry.insert(0, path) + self.path_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.path_button = ctk.CTkButton(self.frame, width=30, text="...", command=lambda: self.browse_for_path(self.path_entry)) + self.path_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) + + self.caption_label = ctk.CTkLabel(self.frame, text="Initial Caption", width=100) + self.caption_label.grid(row=2, column=0, sticky="w", padx=5, pady=5) + self.caption_entry = ctk.CTkEntry(self.frame, width=200) + self.caption_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + + self.prefix_label = ctk.CTkLabel(self.frame, text="Caption Prefix", width=100) + self.prefix_label.grid(row=3, column=0, sticky="w", padx=5, pady=5) + self.prefix_entry = ctk.CTkEntry(self.frame, width=200) + self.prefix_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) + + self.postfix_label = ctk.CTkLabel(self.frame, text="Caption Postfix", width=100) + self.postfix_label.grid(row=4, column=0, sticky="w", padx=5, pady=5) + self.postfix_entry = ctk.CTkEntry(self.frame, width=200) + self.postfix_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + + self.mode_label = ctk.CTkLabel(self.frame, text="Mode", width=100) + self.mode_label.grid(row=5, column=0, sticky="w", padx=5, pady=5) + self.mode_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.mode_var, values=self.modes, dynamic_resizing=False, width=200) + self.mode_dropdown.grid(row=5, column=1, sticky="w", padx=5, pady=5) + + self.include_subdirectories_label = ctk.CTkLabel(self.frame, text="Include subfolders", width=100) + self.include_subdirectories_label.grid(row=6, column=0, sticky="w", padx=5, pady=5) + self.include_subdirectories_var = ctk.BooleanVar(self, parent_include_subdirectories) + self.include_subdirectories_switch = ctk.CTkSwitch(self.frame, text="", variable=self.include_subdirectories_var) + self.include_subdirectories_switch.grid(row=6, column=1, sticky="w", padx=5, pady=5) + + self.progress_label = ctk.CTkLabel(self.frame, text="Progress: 0/0", width=100) + self.progress_label.grid(row=7, column=0, sticky="w", padx=5, pady=5) + self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) + self.progress.grid(row=7, column=1, sticky="w", padx=5, pady=5) + + self.create_captions_button = ctk.CTkButton(self.frame, text="Create Captions", width=310, command=self.create_captions) + self.create_captions_button.grid(row=8, column=0, columnspan=2, sticky="w", padx=5, pady=5) + + self.frame.pack(fill="both", expand=True) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def browse_for_path(self, entry_box): + # get the path from the user + path = filedialog.askdirectory() + # set the path to the entry box + # delete entry box text + entry_box.focus_set() + entry_box.delete(0, filedialog.END) + entry_box.insert(0, path) + self.focus_set() + + def set_progress(self, value, max_value): + progress = value / max_value + self.progress.set(progress) + self.progress_label.configure(text=f"{value}/{max_value}") + self.progress.update() + + def create_captions(self): + self.parent.load_captioning_model(self.model_var.get()) + + mode = { + "Replace all captions": "replace", + "Create if absent": "fill", + "Add as new line": "add", + }[self.mode_var.get()] + + self.parent.captioning_model.caption_folder( + sample_dir=self.path_entry.get(), + initial_caption=self.caption_entry.get(), + caption_prefix=self.prefix_entry.get(), + caption_postfix=self.postfix_entry.get(), + mode=mode, + progress_callback=self.set_progress, + include_subdirectories=self.include_subdirectories_var.get(), + ) + self.parent.load_image() + + def destroy(self): + with contextlib.suppress(tk.TclError): + self.grab_release() + + super().destroy() diff --git a/modules/ui/CtkGenerateMasksWindowView.py b/modules/ui/CtkGenerateMasksWindowView.py new file mode 100644 index 000000000..daff0d3d5 --- /dev/null +++ b/modules/ui/CtkGenerateMasksWindowView.py @@ -0,0 +1,151 @@ +import contextlib +import tkinter as tk +from tkinter import filedialog + +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk + + +class GenerateMasksWindow(ctk.CTkToplevel): + def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs): + """ + Window for generating masks for a folder of images + + Parameters: + parent (`Tk`): the parent window + path (`str`): the path to the folder + parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox + """ + super().__init__(parent, *args, **kwargs) + + self.parent = parent + if path is None: + path = "" + + self.mode_var = ctk.StringVar(self, "Create if absent") + self.modes = ["Replace all masks", "Create if absent", "Add to existing", "Subtract from existing", "Blend with existing"] + self.model_var = ctk.StringVar(self, "ClipSeg") + self.models = ["ClipSeg", "Rembg", "Rembg-Human", "Hex Color"] + + self.title("Batch generate masks") + self.geometry("360x430") + self.resizable(True, True) + + self.frame = ctk.CTkFrame(self, width=600, height=300) + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + + self.model_label = ctk.CTkLabel(self.frame, text="Model", width=100) + self.model_label.grid(row=0, column=0, sticky="w", padx=5, pady=5) + self.model_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.model_var, values=self.models, dynamic_resizing=False, width=200) + self.model_dropdown.grid(row=0, column=1, sticky="w", padx=5, pady=5) + + self.path_label = ctk.CTkLabel(self.frame, text="Folder", width=100) + self.path_label.grid(row=1, column=0, sticky="w",padx=5, pady=5) + self.path_entry = ctk.CTkEntry(self.frame, width=150) + self.path_entry.insert(0, path) + self.path_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.path_button = ctk.CTkButton(self.frame, width=30, text="...", command=lambda: self.browse_for_path(self.path_entry)) + self.path_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) + + self.prompt_label = ctk.CTkLabel(self.frame, text="Prompt", width=100) + self.prompt_label.grid(row=2, column=0, sticky="w",padx=5, pady=5) + self.prompt_entry = ctk.CTkEntry(self.frame, width=200) + self.prompt_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + + self.mode_label = ctk.CTkLabel(self.frame, text="Mode", width=100) + self.mode_label.grid(row=3, column=0, sticky="w", padx=5, pady=5) + self.mode_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.mode_var, values=self.modes, dynamic_resizing=False, width=200) + self.mode_dropdown.grid(row=3, column=1, sticky="w", padx=5, pady=5) + + self.threshold_label = ctk.CTkLabel(self.frame, text="Threshold", width=100) + self.threshold_label.grid(row=4, column=0, sticky="w", padx=5, pady=5) + self.threshold_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="0.0 - 1.0") + self.threshold_entry.insert(0, "0.3") + self.threshold_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + + self.smooth_label = ctk.CTkLabel(self.frame, text="Smooth", width=100) + self.smooth_label.grid(row=5, column=0, sticky="w", padx=5, pady=5) + self.smooth_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="5") + self.smooth_entry.insert(0, 5) + self.smooth_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) + + self.expand_label = ctk.CTkLabel(self.frame, text="Expand", width=100) + self.expand_label.grid(row=6, column=0, sticky="w", padx=5, pady=5) + self.expand_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="10") + self.expand_entry.insert(0, 10) + self.expand_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) + + self.alpha_label = ctk.CTkLabel(self.frame, text="Alpha", width=100) + self.alpha_label.grid(row=7, column=0, sticky="w", padx=5, pady=5) + self.alpha_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="1") + self.alpha_entry.insert(0, 1) + self.alpha_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) + + self.include_subdirectories_label = ctk.CTkLabel(self.frame, text="Include subfolders", width=100) + self.include_subdirectories_label.grid(row=8, column=0, sticky="w", padx=5, pady=5) + self.include_subdirectories_var = ctk.BooleanVar(self, parent_include_subdirectories) + self.include_subdirectories_switch = ctk.CTkSwitch(self.frame, text="", variable=self.include_subdirectories_var) + self.include_subdirectories_switch.grid(row=8, column=1, sticky="w", padx=5, pady=5) + + self.progress_label = ctk.CTkLabel(self.frame, text="Progress: 0/0", width=100) + self.progress_label.grid(row=9, column=0, sticky="w", padx=5, pady=5) + self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) + self.progress.grid(row=9, column=1, sticky="w", padx=5, pady=5) + + self.create_masks_button = ctk.CTkButton(self.frame, text="Create Masks", width=310, command=self.create_masks) + self.create_masks_button.grid(row=10, column=0, columnspan=2, sticky="w", padx=5, pady=5) + + self.frame.pack(fill="both", expand=True) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def browse_for_path(self, entry_box): + # get the path from the user + path = filedialog.askdirectory() + # set the path to the entry box + # delete entry box text + entry_box.focus_set() + entry_box.delete(0, filedialog.END) + entry_box.insert(0, path) + self.focus_set() + + def set_progress(self, value, max_value): + progress = value / max_value + self.progress.set(progress) + self.progress_label.configure(text=f"{value}/{max_value}") + self.progress.update() + + def create_masks(self): + self.parent.load_masking_model(self.model_var.get()) + + mode = { + "Replace all masks": "replace", + "Create if absent": "fill", + "Add to existing": "add", + "Subtract from existing": "subtract", + "Blend with existing": "blend", + }[self.mode_var.get()] + + self.parent.masking_model.mask_folder( + sample_dir=self.path_entry.get(), + prompts=[self.prompt_entry.get()], + mode=mode, + alpha=float(self.alpha_entry.get()), + threshold=float(self.threshold_entry.get()), + smooth_pixels=int(self.smooth_entry.get()), + expand_pixels=int(self.expand_entry.get()), + progress_callback=self.set_progress, + include_subdirectories=self.include_subdirectories_var.get(), + ) + self.parent.load_image() + + def destroy(self): + with contextlib.suppress(tk.TclError): + self.grab_release() + + super().destroy() diff --git a/modules/ui/CtkLoraTabView.py b/modules/ui/CtkLoraTabView.py new file mode 100644 index 000000000..1c73d90ce --- /dev/null +++ b/modules/ui/CtkLoraTabView.py @@ -0,0 +1,154 @@ + +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.DataType import DataType +from modules.util.enum.ModelType import PeftType +from modules.util.ui import components +from modules.util.ui.UIState import UIState +from modules.util.ui.validation_helpers import check_range + +import customtkinter as ctk + + +class LoraTab: + def __init__(self, master, train_config: TrainConfig, ui_state: UIState): + super().__init__() + + self.master = master + self.train_config = train_config + self.ui_state = ui_state + + self.scroll_frame = None + self.options_frame = None + + self.refresh_ui() + + def refresh_ui(self): + if self.scroll_frame: + self.scroll_frame.destroy() + self.scroll_frame = ctk.CTkFrame(self.master, fg_color="transparent") + self.scroll_frame.grid(row=0, column=0, sticky="nsew") + + self.scroll_frame.grid_columnconfigure(0, weight=0) + self.scroll_frame.grid_columnconfigure(1, weight=1) + self.scroll_frame.grid_columnconfigure(2, weight=2) + + components.label(self.scroll_frame, 0, 0, "Type", + tooltip="The type of low-parameter finetuning method.") + # This will instantly call self.setup_lora. + components.options_kv(self.scroll_frame, 0, 1, [ + ("LoRA", PeftType.LORA), + ("LoHa", PeftType.LOHA), + ("OFT v2", PeftType.OFT_2), + ], self.ui_state, "peft_type", command=self.setup_lora) + + def setup_lora(self, peft_type: PeftType): + if peft_type == PeftType.LOHA: + name = "LoHa" + elif peft_type == PeftType.OFT_2: + name = "OFT v2" + else: + name = "LoRA" + + if self.options_frame: + self.options_frame.destroy() + self.options_frame = ctk.CTkFrame(self.scroll_frame, fg_color="transparent") + self.options_frame.grid(row=1, column=0, columnspan=3, sticky="nsew") + master = self.options_frame + + master.grid_columnconfigure(0, weight=0, uniform="a") + master.grid_columnconfigure(1, weight=1, uniform="a") + master.grid_columnconfigure(2, minsize=50, uniform="a") + master.grid_columnconfigure(3, weight=0, uniform="a") + master.grid_columnconfigure(4, weight=1, uniform="a") + + # 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.path_entry( + master, 0, 1, self.ui_state, "lora_model_name", + mode="file", path_modifier=components.json_path_modifier + ) + entry.grid(row=0, column=1, columnspan=4) + + + # LoRA decomposition + if peft_type == PeftType.LORA: + components.label(master, 1, 3, "Decompose Weights (DoRA)", + tooltip="Decompose LoRA Weights (aka, DoRA).") + components.switch(master, 1, 4, self.ui_state, "lora_decompose") + + components.label(master, 2, 3, "Use Norm Epsilon (DoRA Only)", + tooltip="Add an epsilon to the norm divison calculation in DoRA. Can aid in training stability, and also acts as regularization.") + components.switch(master, 2, 4, self.ui_state, "lora_decompose_norm_epsilon") + components.label(master, 3, 3, "Apply on output axis (DoRA Only)", + tooltip="Apply the weight decomposition on the output axis instead of the input axis.") + components.switch(master, 3, 4, self.ui_state, "lora_decompose_output_axis") + + # LoRA and LoHA shared settings + if peft_type == PeftType.LORA or peft_type == PeftType.LOHA: + # rank + components.label(master, 1, 0, f"{name} rank", + tooltip=f"The rank parameter used when creating a new {name}") + components.entry(master, 1, 1, self.ui_state, "lora_rank", required=True, extra_validate=check_range(lower=1, message="Rank must be at least 1")) + + # alpha + components.label(master, 2, 0, f"{name} alpha", + tooltip=f"The alpha parameter used when creating a new {name}") + components.entry(master, 2, 1, self.ui_state, "lora_alpha", required=True) + + # Dropout Percentage + components.label(master, 3, 0, "Dropout Probability", + tooltip="Dropout probability. This percentage of model nodes will be randomly ignored at each training step. Helps with overfitting. 0 disables, 1 maximum.") + components.entry(master, 3, 1, self.ui_state, "dropout_probability") + + # weight dtype + components.label(master, 4, 0, f"{name} Weight Data Type", + tooltip=f"The {name} weight data type used for training. This can reduce memory consumption, but reduces precision") + components.options_kv(master, 4, 1, [ + ("float32", DataType.FLOAT_32), + ("bfloat16", DataType.BFLOAT_16), + ], self.ui_state, "lora_weight_dtype") + + # For use with additional embeddings. + components.label(master, 5, 0, "Bundle Embeddings", + tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") + components.switch(master, 5, 1, self.ui_state, "bundle_additional_embeddings") + + # OFTv2 + elif peft_type == PeftType.OFT_2: + # Block Size + components.label(master, 1, 0, f"{name} Block Size", + tooltip=f"The block size parameter used when creating a new {name}") + components.entry(master, 1, 1, self.ui_state, "oft_block_size", required=True) + + # COFT + components.label(master, 1, 3, "Constrained OFT (COFT)", + tooltip="Use the constrained variant of OFT. This constrains the learned rotation to stay very close to the identity matrix, limiting adaptation to only small changes. This improves training stability, helps prevent overfitting on small datasets, and better preserves the base models original knowledge but it may lack expressiveness for tasks requiring substantial adaptation and introduces an additional hyperparameter (COFT Epsilon) that needs tuning.") + components.switch(master, 1, 4, self.ui_state, "oft_coft") + + components.label(master, 2, 3, "COFT Epsilon", + tooltip="The control strength of COFT. Only has an effect if COFT is enabled.") + components.entry(master, 2, 4, self.ui_state, "coft_eps") + + # Block Share + components.label(master, 3, 3, "Block Share", + tooltip="Share the OFT parameters between blocks. A single rotation matrix is shared across all blocks within a layer, drastically cutting the number of trainable parameters and yielding very compact adapter files, potentially improving generalization but at the cost of significant expressiveness, which can lead to underfitting on more complex or diverse tasks.") + components.switch(master, 3, 4, self.ui_state, "oft_block_share") + + # Dropout Percentage + components.label(master, 2, 0, "Dropout Probability", + tooltip="Dropout probability. This percentage of the rotated adapter nodes that will be randomly restored to the base model initial statue. Helps with overfitting. 0 disables, 1 maximum.") + components.entry(master, 2, 1, self.ui_state, "dropout_probability") + + # OFT weight dtype + components.label(master, 3, 0, f"{name} Weight Data Type", + tooltip=f"The {name} weight data type used for training. This can reduce memory consumption, but reduces precision") + components.options_kv(master, 3, 1, [ + ("float32", DataType.FLOAT_32), + ("bfloat16", DataType.BFLOAT_16), + ], self.ui_state, "lora_weight_dtype") + + # For use with additional embeddings. + components.label(master, 4, 0, "Bundle Embeddings", + tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") + components.switch(master, 4, 1, self.ui_state, "bundle_additional_embeddings") diff --git a/modules/ui/CtkModelTabView.py b/modules/ui/CtkModelTabView.py new file mode 100644 index 000000000..ff17ea3ba --- /dev/null +++ b/modules/ui/CtkModelTabView.py @@ -0,0 +1,688 @@ + +from modules.util import create +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ConfigPart import ConfigPart +from modules.util.enum.DataType import DataType +from modules.util.enum.ModelFormat import ModelFormat +from modules.util.enum.PathIOType import PathIOType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.ui import components +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class ModelTab: + + def __init__(self, master, train_config: TrainConfig, ui_state: UIState): + super().__init__() + + self.master = master + self.train_config = train_config + self.ui_state = ui_state + + master.grid_rowconfigure(0, weight=1) + master.grid_columnconfigure(0, weight=1) + + self.scroll_frame = None + + self.refresh_ui() + + def refresh_ui(self): + if self.scroll_frame: + self.scroll_frame.destroy() + + self.scroll_frame = ctk.CTkScrollableFrame(self.master, fg_color="transparent") + self.scroll_frame.grid(row=0, column=0, sticky="nsew") + self.scroll_frame.grid_columnconfigure(0, weight=1) + + base_frame = ctk.CTkFrame(master=self.scroll_frame, corner_radius=5) + base_frame.grid(row=0, column=0, padx=5, pady=5, sticky="nsew") + + base_frame.grid_columnconfigure(0, weight=0) + base_frame.grid_columnconfigure(1, weight=10)#, minsize=500) + base_frame.grid_columnconfigure(2, minsize=50) + base_frame.grid_columnconfigure(3, weight=0) + base_frame.grid_columnconfigure(4, weight=1) + + if self.train_config.model_type.is_stable_diffusion(): #TODO simplify + self.__setup_stable_diffusion_ui(base_frame) + if self.train_config.model_type.is_stable_diffusion_3(): + self.__setup_stable_diffusion_3_ui(base_frame) + elif self.train_config.model_type.is_stable_diffusion_xl(): + self.__setup_stable_diffusion_xl_ui(base_frame) + elif self.train_config.model_type.is_wuerstchen(): + self.__setup_wuerstchen_ui(base_frame) + elif self.train_config.model_type.is_pixart(): + self.__setup_pixart_alpha_ui(base_frame) + elif self.train_config.model_type.is_flux_1(): + self.__setup_flux_ui(base_frame) + elif self.train_config.model_type.is_flux_2(): + self.__setup_flux_2_ui(base_frame) + elif self.train_config.model_type.is_z_image(): + self.__setup_z_image_ui(base_frame) + elif self.train_config.model_type.is_chroma(): + self.__setup_chroma_ui(base_frame) + elif self.train_config.model_type.is_qwen(): + self.__setup_qwen_ui(base_frame) + elif self.train_config.model_type.is_sana(): + self.__setup_sana_ui(base_frame) + elif self.train_config.model_type.is_hunyuan_video(): + self.__setup_hunyuan_video_ui(base_frame) + elif self.train_config.model_type.is_hi_dream(): + self.__setup_hi_dream_ui(base_frame) + elif self.train_config.model_type.is_ernie(): + self.__setup_ernie_ui(base_frame) + + def __setup_stable_diffusion_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_unet=True, + has_text_encoder=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method in [ + TrainingMethod.FINE_TUNE, + TrainingMethod.FINE_TUNE_VAE, + ], + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_stable_diffusion_3_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + has_text_encoder_1=True, + has_text_encoder_2=True, + has_text_encoder_3=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_flux_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + allow_override_transformer=True, + has_text_encoder_1=True, + has_text_encoder_2=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_flux_2_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + allow_override_transformer=True, + has_text_encoder_1=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_z_image_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + allow_override_transformer=True, + has_text_encoder_1=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_ernie_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + allow_override_transformer=True, + has_text_encoder_1=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_chroma_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + allow_override_transformer=True, + has_text_encoder_1=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_qwen_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + allow_override_transformer=True, + has_text_encoder_1=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_stable_diffusion_xl_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_unet=True, + has_text_encoder_1=True, + has_text_encoder_2=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_wuerstchen_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_prior=True, + allow_override_prior=self.train_config.model_type.is_stable_cascade(), + has_text_encoder=True, + ) + row = self.__create_effnet_encoder_components(frame, row) + row = self.__create_decoder_components(frame, row, self.train_config.model_type.is_wuerstchen_v2()) + row = self.__create_output_components( + frame, + row, + allow_safetensors=self.train_config.training_method != TrainingMethod.FINE_TUNE + or self.train_config.model_type.is_stable_cascade(), + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_pixart_alpha_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + has_text_encoder=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_sana_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + has_text_encoder=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=self.train_config.training_method != TrainingMethod.FINE_TUNE, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_hunyuan_video_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + allow_override_transformer=True, + has_text_encoder_1=True, + has_text_encoder_2=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_hi_dream_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + has_text_encoder_1=True, + has_text_encoder_2=True, + has_text_encoder_3=True, + has_text_encoder_4=True, + allow_override_text_encoder_4=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __create_dtype_options(self, include_gguf: bool=False, include_a8: bool=False) -> list[tuple[str, DataType]]: + options = [ + ("float32", DataType.FLOAT_32), + ("bfloat16", DataType.BFLOAT_16), + ("float16", DataType.FLOAT_16), + ("float8 (W8)", DataType.FLOAT_8), + # ("int8", DataType.INT_8), # TODO: reactivate when the int8 implementation is fixed in bitsandbytes: https://github.com/bitsandbytes-foundation/bitsandbytes/issues/1332 + ("nfloat4", DataType.NFLOAT_4), + ] + if include_a8: + options += [ + ("float W8A8", DataType.FLOAT_W8A8), + ("int W8A8", DataType.INT_W8A8), + ] + + if include_gguf: + options.append(("GGUF", DataType.GGUF)) + if include_a8: + options += [ + ("GGUF A8 float", DataType.GGUF_A8_FLOAT), + ("GGUF A8 int", DataType.GGUF_A8_INT), + ] + + return options + + def __create_base_dtype_components(self, frame, row: int) -> int: + # huggingface token + components.label(frame, row, 0, "Hugging Face Token", + tooltip="Enter your Hugging Face access token if you have used a protected Hugging Face repository below.\nThis value is stored separately, not saved to your configuration file. " + "Go to https://huggingface.co/settings/tokens to create an access token.", + wide_tooltip=True) + components.entry(frame, row, 1, self.ui_state, "secrets.huggingface_token") + + row += 1 + + # base model + components.label(frame, row, 0, "Base Model", + tooltip="Filename, directory or Hugging Face repository of the base model") + components.path_entry( + frame, row, 1, self.ui_state, "base_model_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # compile + components.label(frame, row, 3, "Compile transformer blocks", + tooltip="Uses torch.compile and Triton to significantly speed up training. Only applies to transformer/unet. Disable in case of compatibility issues.") + components.switch(frame, row, 4, self.ui_state, "compile") + + row += 1 + + return row + + def __create_base_components( + self, + frame, + row: int, + has_unet: bool = False, + has_prior: bool = False, + allow_override_prior: bool = False, + has_transformer: bool = False, + allow_override_transformer: bool = False, + allow_override_text_encoder_4: bool = False, + has_text_encoder: bool = False, + has_text_encoder_1: bool = False, + has_text_encoder_2: bool = False, + has_text_encoder_3: bool = False, + has_text_encoder_4: bool = False, + has_vae: bool = False, + ) -> int: + if has_unet: + # unet weight dtype + components.label(frame, row, 3, "UNet Data Type", + tooltip="The unet weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(include_a8=True), + self.ui_state, "unet.weight_dtype") + + row += 1 + + if has_prior: + if allow_override_prior: + # prior model + components.label(frame, row, 0, "Prior Model", + tooltip="Filename, directory or Hugging Face repository of the prior model") + components.path_entry( + frame, row, 1, self.ui_state, "prior.model_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # prior weight dtype + components.label(frame, row, 3, "Prior Data Type", + tooltip="The prior weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "prior.weight_dtype") + + row += 1 + + if has_transformer: + if allow_override_transformer: + # 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.path_entry( + frame, row, 1, self.ui_state, "transformer.model_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # transformer weight dtype + components.label(frame, row, 3, "Transformer Data Type", + tooltip="The transformer weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(include_gguf=True, include_a8=True), + self.ui_state, "transformer.weight_dtype") + + row += 1 + + cls = create.get_model_setup_class(self.train_config.model_type, self.train_config.training_method) + presets = cls.LAYER_PRESETS if cls is not None else {"full": []} + + components.label(frame, row, 0, "Quantization") + components.layer_filter_entry(frame, row, 1, self.ui_state, + preset_var_name="quantization.layer_filter_preset", presets=presets, + preset_label="Quantization Layer Filter", + preset_tooltip="Select a preset defining which layers to quantize. Quantization of certain layers can decrease model quality. Only applies to the transformer/unet", + entry_var_name="quantization.layer_filter", + entry_tooltip="Comma-separated list of layers to quantize. Regular expressions (if toggled) are supported. Any model layer with a matching name will be quantized", + regex_var_name="quantization.layer_filter_regex", + regex_tooltip="If enabled, layer filter patterns are interpreted as regular expressions. Otherwise, simple substring matching is used.", + frame_color="transparent", + ) + + # SVDQuant - create vertical grids to match the size of layer_filter_entry + svd_label_frame = ctk.CTkFrame(frame, fg_color="transparent") + svd_label_frame.grid(row=row, column=3, sticky="nsew") + svd_entry_frame = ctk.CTkFrame(frame, fg_color="transparent") + svd_entry_frame.grid(row=row, column=4, sticky="nsew") + components.label(svd_label_frame, 0, 0, "SVDQuant", + tooltip="What datatype to use for SVDQuant weights decomposition.") + components.options_kv(svd_entry_frame, 0, 0, [("disabled", DataType.NONE), ("float32", DataType.FLOAT_32), ("bfloat16", DataType.BFLOAT_16)], + self.ui_state, "quantization.svd_dtype") + components.label(svd_label_frame, 1, 0, "SVDQuant Rank", + tooltip="Rank for SVDQuant weights decomposition") + components.entry(svd_entry_frame, 1, 0, self.ui_state, "quantization.svd_rank") + row += 1 + + + if has_text_encoder: + # text encoder weight dtype + components.label(frame, row, 3, "Text Encoder Data Type", + tooltip="The text encoder weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "text_encoder.weight_dtype") + + row += 1 + + if has_text_encoder_1: + # text encoder 1 weight dtype + components.label(frame, row, 3, "Text Encoder 1 Data Type", + tooltip="The text encoder 1 weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "text_encoder.weight_dtype") + + row += 1 + + if has_text_encoder_2: + # text encoder 2 weight dtype + components.label(frame, row, 3, "Text Encoder 2 Data Type", + tooltip="The text encoder 2 weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "text_encoder_2.weight_dtype") + + row += 1 + + if has_text_encoder_3: + # text encoder 3 weight dtype + components.label(frame, row, 3, "Text Encoder 3 Data Type", + tooltip="The text encoder 3 weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "text_encoder_3.weight_dtype") + + row += 1 + + if has_text_encoder_4: + if allow_override_text_encoder_4: + # 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.path_entry( + frame, row, 1, self.ui_state, "text_encoder_4.model_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # text encoder 4 weight dtype + components.label(frame, row, 3, "Text Encoder 4 Data Type", + tooltip="The text encoder 4 weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "text_encoder_4.weight_dtype") + + row += 1 + + 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.path_entry( + frame, row, 1, self.ui_state, "vae.model_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # vae weight dtype + components.label(frame, row, 3, "VAE Data Type", + tooltip="The vae weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "vae.weight_dtype") + + row += 1 + + return row + + 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.path_entry( + frame, row, 1, self.ui_state, "effnet_encoder.model_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # effnet encoder weight dtype + components.label(frame, row, 3, "Effnet Encoder Data Type", + tooltip="The effnet encoder weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "effnet_encoder.weight_dtype") + + row += 1 + + return row + + def __create_decoder_components( + self, + frame, + row: int, + has_text_encoder: bool, + ) -> int: + # decoder model + components.label(frame, row, 0, "Decoder Model", + tooltip="Filename, directory or Hugging Face repository of the decoder model") + components.path_entry( + frame, row, 1, self.ui_state, "decoder.model_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # decoder weight dtype + components.label(frame, row, 3, "Decoder Data Type", + tooltip="The decoder weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "decoder.weight_dtype") + + row += 1 + + if has_text_encoder: + # decoder text encoder weight dtype + components.label(frame, row, 3, "Decoder Text Encoder Data Type", + tooltip="The decoder text encoder weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "decoder_text_encoder.weight_dtype") + + row += 1 + + # decoder vqgan weight dtype + components.label(frame, row, 3, "Decoder VQGAN Data Type", + tooltip="The decoder vqgan weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "decoder_vqgan.weight_dtype") + + row += 1 + + return row + + def __create_output_components( + self, + frame, + row: int, + allow_safetensors: bool = False, + allow_diffusers: bool = False, + allow_legacy_safetensors: bool = False, + allow_comfy: bool = False, + ) -> int: + # output model destination + components.label(frame, row, 0, "Model Output Destination", + tooltip="Filename or directory where the output model is saved") + components.path_entry( + frame, row, 1, self.ui_state, "output_model_destination", + mode="file", + io_type=PathIOType.MODEL, + ) + + # output data type + components.label(frame, row, 3, "Output Data Type", + tooltip="Precision to use when saving the output model") + components.options_kv(frame, row, 4, [ + ("float16", DataType.FLOAT_16), + ("float32", DataType.FLOAT_32), + ("bfloat16", DataType.BFLOAT_16), + ("float8", DataType.FLOAT_8), + ("nfloat4", DataType.NFLOAT_4), + ], self.ui_state, "output_dtype") + + row += 1 + + # output format + formats = [] + if allow_safetensors: + formats.append(("Safetensors", ModelFormat.SAFETENSORS)) + if allow_diffusers: + formats.append(("Diffusers", ModelFormat.DIFFUSERS)) + # if allow_legacy_safetensors: + # formats.append(("Legacy Safetensors", ModelFormat.LEGACY_SAFETENSORS)) + if allow_comfy: + formats.append(("Comfy LoRA", ModelFormat.COMFY_LORA)) + + 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") + + # include config + components.label(frame, row, 3, "Include Config", + tooltip="Include the training configuration in the final model. Only supported for safetensors files. " + "None: No config is included. " + "Settings: All training settings are included. " + "All: All settings, including the samples and concepts are included.") + components.options_kv(frame, row, 4, [ + ("None", ConfigPart.NONE), + ("Settings", ConfigPart.SETTINGS), + ("All", ConfigPart.ALL), + ], self.ui_state, "include_train_config") + + row += 1 + + return row diff --git a/modules/ui/CtkMuonAdamWindowView.py b/modules/ui/CtkMuonAdamWindowView.py new file mode 100644 index 000000000..5879ab432 --- /dev/null +++ b/modules/ui/CtkMuonAdamWindowView.py @@ -0,0 +1,107 @@ +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.Optimizer import Optimizer +from modules.util.optimizer_util import OPTIMIZER_DEFAULT_PARAMETERS +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + +MUON_AUX_ADAM_DEFAULTS = { + "beta1": 0.9, + "beta2": 0.999, + "eps": 1e-8, + "weight_decay": 0.0, +} + +class MuonAdamWindow(ctk.CTkToplevel): + def __init__( + self, + parent, + train_config: TrainConfig, + ui_state: UIState, + parent_optimizer_type: Optimizer, + *args, **kwargs, + ): + super().__init__(parent, *args, **kwargs) + + self.parent = parent + self.train_config = train_config + self.adam_ui_state = ui_state + self.parent_optimizer_type = parent_optimizer_type + + if self.parent_optimizer_type == Optimizer.MUON: + self.title("Muon's Auxiliary AdamW Settings") + self.adam_params_def = MUON_AUX_ADAM_DEFAULTS + else: + self.title("Muon_adv's Auxiliary AdamW_adv Settings") + self.adam_params_def = OPTIMIZER_DEFAULT_PARAMETERS[Optimizer.ADAMW_ADV] + + self.geometry("800x500") + self.resizable(True, True) + + self.grid_rowconfigure(0, weight=1) + self.grid_rowconfigure(1, weight=0) + self.grid_columnconfigure(0, weight=1) + + self.frame = ctk.CTkScrollableFrame(self, fg_color="transparent") + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + + self.frame.grid_columnconfigure(0, weight=0) + self.frame.grid_columnconfigure(1, weight=1) + self.frame.grid_columnconfigure(2, minsize=50) + self.frame.grid_columnconfigure(3, weight=0) + self.frame.grid_columnconfigure(4, weight=1) + + components.button(self, 1, 0, "ok", command=self.destroy) + self.create_adam_params_ui(self.frame) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + def create_adam_params_ui(self, master): + # This is a large map, copied from OptimizerParamsWindow for simplicity. + # @formatter:off + KEY_DETAIL_MAP = { + 'alpha': {'title': 'Alpha', 'tooltip': 'Smoothing parameter for RMSprop and others.', 'type': 'float'}, + 'beta1': {'title': 'Beta1', 'tooltip': 'optimizer_momentum term.', 'type': 'float'}, + 'beta2': {'title': 'Beta2', 'tooltip': 'Coefficients for computing running averages of gradient.', 'type': 'float'}, + 'eps': {'title': 'EPS', 'tooltip': 'A small value to prevent division by zero.', 'type': 'float'}, + 'stochastic_rounding': {'title': 'Stochastic Rounding', 'tooltip': 'Stochastic rounding for weight updates. Improves quality when using bfloat16 weights.', 'type': 'bool'}, + 'use_bias_correction': {'title': 'Bias Correction', 'tooltip': 'Turn on Adam\'s bias correction.', 'type': 'bool'}, + 'weight_decay': {'title': 'Weight Decay', 'tooltip': 'Regularization to prevent overfitting.', 'type': 'float'}, + 'use_orthograd': {'title': 'use_orthograd', 'tooltip': 'Use orthograd method', 'type': 'bool'}, + 'nnmf_factor': {'title': 'Factored Optimizer', 'tooltip': 'Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. This is an approximation which may slightly alter training dynamics.', 'type': 'bool'}, + 'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization.', 'type': 'bool'}, + 'use_atan2': {'title': 'Atan2 Scaling', 'tooltip': 'A robust replacement for eps, which also incorporates gradient clipping, bounding and stabilizing the optimizer updates.', 'type': 'bool'}, + 'use_AdEMAMix': {'title': 'AdEMAMix EMA', 'tooltip': 'Adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, and accelerate the training.', 'type': 'bool'}, + 'beta3_ema': {'title': 'Beta3 EMA', 'tooltip': 'Coefficient for slow-moving EMA of AdEMAMix.', 'type': 'float'}, + 'Simplified_AdEMAMix': {'title': 'Simplified AdEMAMix', 'tooltip': "Enables a simplified, single-EMA variant of AdEMAMix. Instead of blending two moving averages (fast and slow momentum), this version combines the raw current gradient (controlled by 'Grad α') directly with a single theory-based momentum. This makes the optimizer highly responsive to recent gradient information, which can accelerate training in all batch size scenarios when tuned correctly.", 'type': 'bool'}, + 'alpha_grad': {'title': 'Grad α', 'tooltip': 'Controls the mixing coefficient between raw gradients and momentum gradients in Simplified AdEMAMix. Higher values (e.g., 10-100) emphasize recent gradients, suitable for small batch sizes to reduce noise. Lower values (e.g., 0-1) emphasize historical gradients, suitable for large batch sizes for stability. Setting to 0 uses only momentum gradients without raw gradient contribution.', 'type': 'float'}, + 'kourkoutas_beta': {'title': 'Kourkoutas Beta', 'tooltip': 'Enables a layer-wise dynamic β₂ adaptation. This feature makes the optimizer more responsive to "spiky" gradients by lowering β₂ during periods of high variance, and more stable during calm periods by raising β₂ towards its maximum. It can significantly improve training stability and final loss.', 'type': 'bool'}, + } + # @formatter:on + + adam_params = self.adam_params_def + + for index, key in enumerate(adam_params.keys()): + if key not in KEY_DETAIL_MAP: + continue + + arg_info = KEY_DETAIL_MAP[key] + + title = arg_info['title'] + tooltip = arg_info['tooltip'] + param_type = arg_info['type'] + + row = index // 2 + col = 3 * (index % 2) + + components.label(master, row, col, title, tooltip=tooltip) + + if param_type != 'bool': + components.entry(master, row, col + 1, self.adam_ui_state, key) + else: + components.switch(master, row, col + 1, self.adam_ui_state, key) diff --git a/modules/ui/CtkOffloadingWindowView.py b/modules/ui/CtkOffloadingWindowView.py new file mode 100644 index 000000000..54035e121 --- /dev/null +++ b/modules/ui/CtkOffloadingWindowView.py @@ -0,0 +1,75 @@ +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.GradientCheckpointingMethod import ( + GradientCheckpointingMethod, +) +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class OffloadingWindow(ctk.CTkToplevel): + def __init__( + self, + parent, + config: TrainConfig, + ui_state: UIState, + *args, **kwargs, + ): + super().__init__(parent, *args, **kwargs) + + self.config = config + self.ui_state = ui_state + self.image_preview_file_index = 0 + self.ax = None + self.canvas = None + + self.title("Offloading") + self.geometry("800x400") + self.resizable(True, True) + + self.grid_rowconfigure(0, weight=1) + self.grid_columnconfigure(0, weight=1) + + frame = self.__content_frame(self) + frame.grid(row=0, column=0, sticky='nsew') + components.button(self, 1, 0, "ok", self.__ok) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def __content_frame(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=1) + frame.grid_columnconfigure(1, weight=1) + + # timestep distribution + components.label(frame, 0, 0, "Gradient checkpointing", + tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") + components.options(frame, 0, 1, [str(x) for x in list(GradientCheckpointingMethod)], self.ui_state, + "gradient_checkpointing") + + # gradient checkpointing layer offloading + components.label(frame, 1, 0, "Async Offloading", + tooltip="Enables Asynchronous offloading.") + components.switch(frame, 1, 1, self.ui_state, "enable_async_offloading") + + # gradient checkpointing layer offloading + components.label(frame, 2, 0, "Offload Activations", + tooltip="Enables Activation Offloading") + components.switch(frame, 2, 1, self.ui_state, "enable_activation_offloading") + + # gradient checkpointing layer offloading + components.label(frame, 3, 0, "Layer offload fraction", + tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") + components.entry(frame, 3, 1, self.ui_state, "layer_offload_fraction") + + frame.pack(fill="both", expand=1) + return frame + + def __ok(self): + self.destroy() diff --git a/modules/ui/CtkOptimizerParamsWindowView.py b/modules/ui/CtkOptimizerParamsWindowView.py new file mode 100644 index 000000000..16063c26c --- /dev/null +++ b/modules/ui/CtkOptimizerParamsWindowView.py @@ -0,0 +1,288 @@ +import contextlib +from tkinter import TclError + +from modules.ui.MuonAdamWindow import MUON_AUX_ADAM_DEFAULTS, MuonAdamWindow +from modules.util.config.TrainConfig import TrainConfig, TrainOptimizerConfig +from modules.util.enum.Optimizer import Optimizer +from modules.util.optimizer_util import ( + OPTIMIZER_DEFAULT_PARAMETERS, + change_optimizer, + load_optimizer_defaults, + update_optimizer_config, +) +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class OptimizerParamsWindow(ctk.CTkToplevel): + def __init__( + self, + parent, + train_config: TrainConfig, + ui_state, + *args, **kwargs, + ): + super().__init__(parent, *args, **kwargs) + + self.parent = parent + self.train_config = train_config + self.ui_state = ui_state + self.optimizer_ui_state = ui_state.get_var("optimizer") + self.protocol("WM_DELETE_WINDOW", self.on_window_close) + self.muon_adam_button = None + + self.title("Optimizer Settings") + self.geometry("800x500") + self.resizable(True, True) + + self.grid_rowconfigure(0, weight=1) + self.grid_rowconfigure(1, weight=0) + self.grid_columnconfigure(0, weight=1) + + self.frame = ctk.CTkScrollableFrame(self, fg_color="transparent") + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + + self.frame.grid_columnconfigure(0, weight=0) + self.frame.grid_columnconfigure(1, weight=1) + self.frame.grid_columnconfigure(2, minsize=50) + self.frame.grid_columnconfigure(3, weight=0) + self.frame.grid_columnconfigure(4, weight=1) + + components.button(self, 1, 0, "ok", command=self.on_window_close) + self.main_frame(self.frame) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def main_frame(self, master): + # Optimizer + components.label(master, 0, 0, "Optimizer", + tooltip="The type of optimizer") + + # Create the optimizer dropdown menu and set the command + components.options(master, 0, 1, [str(x) for x in list(Optimizer)], self.optimizer_ui_state, "optimizer", + command=self.on_optimizer_change) + + # Defaults Button + components.label(master, 0, 3, "Optimizer Defaults", + tooltip="Load default settings for the selected optimizer") + components.button(self.frame, 0, 4, "Load Defaults", self.load_defaults, + tooltip="Load default settings for the selected optimizer") + + self.create_dynamic_ui(master) + + def clear_dynamic_ui(self, master): + with contextlib.suppress(TclError): + for widget in master.winfo_children(): + grid_info = widget.grid_info() + if int(grid_info["row"]) >= 1: + widget.destroy() + + def create_dynamic_ui( + self, + master, + ): + + # Lookup for the title and tooltip for a key + # @formatter:off + KEY_DETAIL_MAP = { + 'adam_w_mode': {'title': 'Adam W Mode', 'tooltip': 'Whether to use weight decay correction for Adam optimizer.', 'type': 'bool'}, + 'alpha': {'title': 'Alpha', 'tooltip': 'Smoothing parameter for RMSprop and others.', 'type': 'float'}, + 'amsgrad': {'title': 'AMSGrad', 'tooltip': 'Whether to use the AMSGrad variant for Adam.', 'type': 'bool'}, + 'beta1': {'title': 'Beta1', 'tooltip': 'optimizer_momentum term.', 'type': 'float'}, + 'beta2': {'title': 'Beta2', 'tooltip': 'Coefficients for computing running averages of gradient.', 'type': 'float'}, + 'beta3': {'title': 'Beta3', 'tooltip': 'Coefficient for computing the Prodigy stepsize.', 'type': 'float'}, + 'bias_correction': {'title': 'Bias Correction', 'tooltip': 'Whether to use bias correction in optimization algorithms like Adam.', 'type': 'bool'}, + 'block_wise': {'title': 'Block Wise', 'tooltip': 'Whether to perform block-wise model update.', 'type': 'bool'}, + 'capturable': {'title': 'Capturable', 'tooltip': 'Whether some property of the optimizer can be captured.', 'type': 'bool'}, + 'centered': {'title': 'Centered', 'tooltip': 'Whether to center the gradient before scaling. Great for stabilizing the training process.', 'type': 'bool'}, + 'clip_threshold': {'title': 'Clip Threshold', 'tooltip': 'Clipping value for gradients.', 'type': 'float'}, + 'd0': {'title': 'Initial D', 'tooltip': 'Initial D estimate for D-adaptation.', 'type': 'float'}, + 'd_coef': {'title': 'D Coefficient', 'tooltip': 'Coefficient in the expression for the estimate of d.', 'type': 'float'}, + 'dampening': {'title': 'Dampening', 'tooltip': 'Dampening for optimizer_momentum.', 'type': 'float'}, + 'decay_rate': {'title': 'Decay Rate', 'tooltip': 'Rate of decay for moment estimation.', 'type': 'float'}, + 'decouple': {'title': 'Decouple', 'tooltip': 'Use AdamW style optimizer_decoupled weight decay.', 'type': 'bool'}, + 'differentiable': {'title': 'Differentiable', 'tooltip': 'Whether the optimization function is optimizer_differentiable.', 'type': 'bool'}, + 'eps': {'title': 'EPS', 'tooltip': 'A small value to prevent division by zero.', 'type': 'float'}, + 'eps2': {'title': 'EPS 2', 'tooltip': 'A small value to prevent division by zero.', 'type': 'float'}, + 'foreach': {'title': 'ForEach', 'tooltip': 'Whether to use a foreach implementation if available. This implementation is usually faster.', 'type': 'bool'}, + 'fsdp_in_use': {'title': 'FSDP in Use', 'tooltip': 'Flag for using sharded parameters.', 'type': 'bool'}, + 'fused': {'title': 'Fused', 'tooltip': 'Whether to use a fused implementation if available. This implementation is usually faster and requires less memory.', 'type': 'bool'}, + 'fused_back_pass': {'title': 'Fused Back Pass', 'tooltip': 'Whether to fuse the back propagation pass with the optimizer step. This reduces VRAM usage, but is not compatible with gradient accumulation.', 'type': 'bool'}, + 'growth_rate': {'title': 'Growth Rate', 'tooltip': 'Limit for D estimate growth rate.', 'type': 'float'}, + 'initial_accumulator_value': {'title': 'Initial Accumulator Value', 'tooltip': 'Initial value for Adagrad optimizer.', 'type': 'float'}, + 'initial_accumulator': {'title': 'Initial Accumulator', 'tooltip': 'Sets the starting value for both moment estimates to ensure numerical stability and balanced adaptive updates early in training.', 'type': 'float'}, + 'is_paged': {'title': 'Is Paged', 'tooltip': 'Whether the optimizer\'s internal state should be paged to CPU.', 'type': 'bool'}, + 'log_every': {'title': 'Log Every', 'tooltip': 'Intervals at which logging should occur.', 'type': 'int'}, + 'lr_decay': {'title': 'LR Decay', 'tooltip': 'Rate at which learning rate decreases.', 'type': 'float'}, + 'max_unorm': {'title': 'Max Unorm', 'tooltip': 'Maximum value for gradient clipping by norms.', 'type': 'float'}, + 'maximize': {'title': 'Maximize', 'tooltip': 'Whether to optimizer_maximize the optimization function.', 'type': 'bool'}, + 'min_8bit_size': {'title': 'Min 8bit Size', 'tooltip': 'Minimum tensor size for 8-bit quantization.', 'type': 'int'}, + 'quant_block_size': {'title': 'Quant Block Size', 'tooltip': 'Size of a block of normalized 8-bit quantization data. Larger values increase memory efficiency at the cost of data precision.', 'type': 'int'}, + 'momentum': {'title': 'optimizer_momentum', 'tooltip': 'Factor to accelerate SGD in relevant direction.', 'type': 'float'}, + 'nesterov': {'title': 'Nesterov', 'tooltip': 'Whether to enable Nesterov optimizer_momentum.', 'type': 'bool'}, + 'no_prox': {'title': 'No Prox', 'tooltip': 'Whether to use proximity updates or not.', 'type': 'bool'}, + 'optim_bits': {'title': 'Optim Bits', 'tooltip': 'Number of bits used for optimization.', 'type': 'int'}, + 'percentile_clipping': {'title': 'Percentile Clipping', 'tooltip': 'Gradient clipping based on percentile values.', 'type': 'int'}, + 'relative_step': {'title': 'Relative Step', 'tooltip': 'Whether to use a relative step size.', 'type': 'bool'}, + 'safeguard_warmup': {'title': 'Safeguard Warmup', 'tooltip': 'Avoid issues during warm-up stage.', 'type': 'bool'}, + 'scale_parameter': {'title': 'Scale Parameter', 'tooltip': 'Whether to scale the parameter or not.', 'type': 'bool'}, + 'stochastic_rounding': {'title': 'Stochastic Rounding', 'tooltip': 'Stochastic rounding for weight updates. Improves quality when using bfloat16 weights.', 'type': 'bool'}, + 'use_bias_correction': {'title': 'Bias Correction', 'tooltip': 'Turn on Adam\'s bias correction.', 'type': 'bool'}, + 'use_triton': {'title': 'Use Triton', 'tooltip': 'Whether Triton optimization should be used.', 'type': 'bool'}, + 'warmup_init': {'title': 'Warmup Initialization', 'tooltip': 'Whether to warm-up the optimizer initialization.', 'type': 'bool'}, + 'weight_decay': {'title': 'Weight Decay', 'tooltip': 'Regularization to prevent overfitting.', 'type': 'float'}, + 'weight_lr_power': {'title': 'Weight LR Power', 'tooltip': 'During warmup, the weights in the average will be equal to lr raised to this power. Set to 0 for no weighting.', 'type': 'float'}, + 'decoupled_decay': {'title': 'Decoupled Decay', 'tooltip': 'If set as True, then the optimizer uses decoupled weight decay as in AdamW.', 'type': 'bool'}, + 'fixed_decay': {'title': 'Fixed Decay', 'tooltip': '(When Decoupled Decay is True:) Applies fixed weight decay when True; scales decay with learning rate when False.', 'type': 'bool'}, + 'rectify': {'title': 'Rectify', 'tooltip': 'Perform the rectified update similar to RAdam.', 'type': 'bool'}, + 'degenerated_to_sgd': {'title': 'Degenerated to SGD', 'tooltip': 'Performs SGD update when gradient variance is high.', 'type': 'bool'}, + 'k': {'title': 'K', 'tooltip': 'Number of vector projected per iteration.', 'type': 'int'}, + 'xi': {'title': 'Xi', 'tooltip': 'Term used in vector projections to avoid division by zero.', 'type': 'float'}, + 'n_sma_threshold': {'title': 'N SMA Threshold', 'tooltip': 'Number of SMA threshold.', 'type': 'int'}, + 'ams_bound': {'title': 'AMS Bound', 'tooltip': 'Whether to use the AMSBound variant.', 'type': 'bool'}, + 'r': {'title': 'R', 'tooltip': 'EMA factor.', 'type': 'float'}, + 'adanorm': {'title': 'AdaNorm', 'tooltip': 'Whether to use the AdaNorm variant', 'type': 'bool'}, + 'adam_debias': {'title': 'Adam Debias', 'tooltip': 'Only correct the denominator to avoid inflating step sizes early in training.', 'type': 'bool'}, + 'slice_p': {'title': 'Slice parameters', 'tooltip': 'Reduce memory usage by calculating LR adaptation statistics on only every pth entry of each tensor. For values greater than 1 this is an approximation to standard Prodigy. Values ~11 are reasonable.', 'type': 'int'}, + 'cautious': {'title': 'Cautious', 'tooltip': 'Whether to use the Cautious variant', 'type': 'bool'}, + 'weight_decay_by_lr': {'title': 'weight_decay_by_lr', 'tooltip': 'Automatically adjust weight decay based on lr', 'type': 'bool'}, + 'prodigy_steps': {'title': 'prodigy_steps', 'tooltip': 'Turn off Prodigy after N steps', 'type': 'int'}, + 'use_speed': {'title': 'use_speed', 'tooltip': 'use_speed method', 'type': 'bool'}, + 'split_groups': {'title': 'split_groups', 'tooltip': 'Use split groups when training multiple params(uNet,TE..)', 'type': 'bool'}, + 'split_groups_mean': {'title': 'split_groups_mean', 'tooltip': 'Use mean for split groups', 'type': 'bool'}, + 'factored': {'title': 'factored', 'tooltip': 'Use factored', 'type': 'bool'}, + 'factored_fp32': {'title': 'factored_fp32', 'tooltip': 'Use factored_fp32', 'type': 'bool'}, + 'use_stableadamw': {'title': 'use_stableadamw', 'tooltip': 'Use use_stableadamw for gradient scaling', 'type': 'bool'}, + 'use_cautious': {'title': 'use_cautious', 'tooltip': 'Use cautious method', 'type': 'bool'}, + 'use_grams': {'title': 'use_grams', 'tooltip': 'Use grams method', 'type': 'bool'}, + 'use_adopt': {'title': 'use_adopt', 'tooltip': 'Use adopt method', 'type': 'bool'}, + 'd_limiter': {'title': 'd_limiter', 'tooltip': 'Prevent over-estimated LRs when gradients and EMA are still stabilizing', 'type': 'bool'}, + 'use_schedulefree': {'title': 'use_schedulefree', 'tooltip': 'Use Schedulefree method', 'type': 'bool'}, + 'use_orthograd': {'title': 'use_orthograd', 'tooltip': 'Use orthograd method', 'type': 'bool'}, + 'nnmf_factor': {'title': 'Factored Optimizer', 'tooltip': 'Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. This is an approximation which may slightly alter training dynamics.', 'type': 'bool'}, + 'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization.', 'type': 'bool'}, + 'use_atan2': {'title': 'Atan2 Scaling', 'tooltip': 'A robust replacement for eps, which also incorporates gradient clipping, bounding and stabilizing the optimizer updates.', 'type': 'bool'}, + 'use_AdEMAMix': {'title': 'AdEMAMix EMA', 'tooltip': 'Adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, and accelerate the training.', 'type': 'bool'}, + 'beta3_ema': {'title': 'Beta3 EMA', 'tooltip': 'Coefficient for slow-moving EMA of AdEMAMix.', 'type': 'float'}, + 'beta1_warmup': {'title': 'Beta1 Warmup Steps', 'tooltip': 'Number of warmup steps to gradually increase beta1 from Minimum Beta1 Value to its final value. During warmup, beta1 increases linearly. leave it empty to disable warmup and use constant beta1.', 'type': 'int'}, + 'min_beta1': {'title': 'Minimum Beta1', 'tooltip': 'Starting beta1 value for warmup scheduling. Used only when beta1 warmup is enabled. Lower values allow faster initial adaptation, while higher values provide more smoothing. The final beta1 value is specified in the beta1 parameter.', 'type': 'float'}, + 'Simplified_AdEMAMix': {'title': 'Simplified AdEMAMix', 'tooltip': "Enables a simplified, single-EMA variant of AdEMAMix. Instead of blending two moving averages (fast and slow momentum), this version combines the raw current gradient (controlled by 'Grad α') directly with a single theory-based momentum. This makes the optimizer highly responsive to recent gradient information, which can accelerate training in all batch size scenarios when tuned correctly.", 'type': 'bool'}, + 'alpha_grad': {'title': 'Grad α', 'tooltip': 'Controls the mixing coefficient between raw gradients and momentum gradients in Simplified AdEMAMix. Higher values (e.g., 10-100) emphasize recent gradients, suitable for small batch sizes to reduce noise. Lower values (e.g., 0-1) emphasize historical gradients, suitable for large batch sizes for stability. Setting to 0 uses only momentum gradients without raw gradient contribution.', 'type': 'float'}, + 'kourkoutas_beta': {'title': 'Kourkoutas Beta', 'tooltip': 'Enables a layer-wise dynamic β₂ adaptation. This feature makes the optimizer more responsive to "spiky" gradients by lowering β₂ during periods of high variance, and more stable during calm periods by raising β₂ towards its maximum. It can significantly improve training stability and final loss.', 'type': 'bool'}, + 'schedulefree_c': {'title': 'Schedule free averaging strength', 'tooltip': 'Larger values = more responsive (shorter averaging window); smaller values = smoother (longer window). Set to 0 to disable and use the original Schedule-Free rule. Short small batches (≈6-12); long/large-batch (≈50-200).', 'type': 'float'}, + 'ns_steps': {'title': 'Newton-Schulz Iterations', 'tooltip': 'Controls the number of iterations for update orthogonalization. Higher values improve the updates quality but make each step slower. Lower values are faster per step but may be less effective.', 'type': 'int'}, + 'MuonWithAuxAdam': {'title': 'MuonWithAuxAdam', 'tooltip': 'Whether to use the standard way of Muon. Non-hidden layers fallback to ADAMW, and MUON takes the rest. Note: The auxiliary Adam (ADAMW) is typically only relevant for training "full" LoRA (LoRA for all layers) or full finetune and is irrelevant for most common LoRA use cases.', 'type': 'bool'}, + 'muon_hidden_layers': {'title': 'Hidden Layers', 'tooltip': 'Comma-separated list of hidden layers to train using Muon. Regular expressions (if toggled) are supported. Any model layer with a matching name will be trained using Muon. If None is provided it will default to using automatic way of finding hidden layers.', 'type': 'str'}, + 'muon_adam_regex': {'title': 'Use Regex', 'tooltip': 'Whether to use regular expressions for hidden layers.', 'type': 'bool'}, + 'muon_adam_lr': {'title': 'Auxiliary Adam LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer. If empty, it will use the main learning rate.', 'type': 'float'}, + 'muon_te1_adam_lr': {'title': 'AuxAdam TE1 LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer for the first text encoder. If empty, it will use the Auxiliary Adam LR.', 'type': 'float'}, + 'muon_te2_adam_lr': {'title': 'AuxAdam TE2 LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer for the second text encoder. If empty, it will use the Auxiliary Adam LR.', 'type': 'float'}, + 'rms_rescaling': {'title': 'RMS Rescaling', 'tooltip': 'Muon already scales its updates to approximate and use the same learning rate (LR) as Adam. This option integrates a more accurate method to match the Adam LR, but it is slower.', 'type': 'bool'}, + 'normuon_variant': {'title': 'NorMuon Variant', 'tooltip': 'Enables the NorMuon optimizer variant, which combines Muon orthogonalization with per-neuron adaptive learning rates for better convergence and balanced parameter updates. Costs only one scalar state buffer per parameter group, size few KBs, maintaining high memory efficiency.', 'type': 'bool'}, + 'beta2_normuon': {'title': 'NorMuon Beta2', 'tooltip': 'Exponential decay rate for the neuron-wise second-moment estimator in NorMuon (analogous to Adams beta2). Controls how past squared updates influence current normalization.', 'type': 'float'}, + 'low_rank_ortho': {'title': 'Low-rank Orthogonalization', 'tooltip': 'Use low-rank orthogonalization to accelerate Muon by orthogonalizing only in a low-dimensional subspace, improving speed and noise robustness.', 'type': 'bool'}, + 'ortho_rank': {'title': 'Ortho Rank', 'tooltip': 'Target rank for low-rank orthogonalization. Controls the dimensionality of the subspace used for efficient and noise-robust orthogonalization.', 'type': 'int'}, + 'accelerated_ns': {'title': 'Accelerated Newton-Schulz', 'tooltip': 'Applies an enhanced Newton-Schulz variant that replaces heuristic coefficients with optimal coefficients derived at each step. This improves performance and convergence by reducing the number of required operations.', 'type': 'bool'}, + 'cautious_wd': {'title': 'Cautious Weight Decay', 'tooltip': 'Applies weight decay only to parameter coordinates whose signs align with the optimizer update direction. This preserves the original optimization objective while still benefiting from regularization effects, leading to improved convergence and better final performance.', 'type': 'bool'}, + 'approx_mars': {'title': 'Approx MARS-M', 'tooltip': 'Enables Approximated MARS-M, a variance reduction technique. It uses the previous step\'s gradient to correct the current update, leading to lower losses and improved convergence stability. This requires additional state to store the previous gradient.', 'type': 'bool'}, + 'auto_kappa_p': {'title': 'Auto Lion-K', 'tooltip': 'Automatically determines the optimal P-value based on layer dimensions. Uses p=2.0 (Spherical) for 4D (Conv) tensors for stability and rotational invariance, and p=1.0 (Sign) for 2D (Linear) tensors for sparsity. Overrides the manual P-value. Recommend for unet models.', 'type': 'bool'}, + 'compile': {'title': 'Compiled Optimizer', 'tooltip': 'Enables PyTorch compilation for the optimizer internal step logic. This is intended to improve performance by allowing PyTorch to fuse operations and optimize the computational graph.', 'type': 'bool'}, + } + # @formatter:on + + if not self.winfo_exists(): # check if this window isn't open + return + + selected_optimizer = self.train_config.optimizer.optimizer + + # Extract the keys for the selected optimizer + for index, key in enumerate(OPTIMIZER_DEFAULT_PARAMETERS[selected_optimizer].keys()): + if key not in KEY_DETAIL_MAP: + continue + arg_info = KEY_DETAIL_MAP[key] + + title = arg_info['title'] + tooltip = arg_info['tooltip'] + type = arg_info['type'] + + row = (index // 2) + 1 + col = 3 * (index % 2) + + components.label(master, row, col, title, tooltip=tooltip) + + if key == 'MuonWithAuxAdam': + frame = ctk.CTkFrame(master, fg_color="transparent") + frame.grid(row=row, column=col + 1, columnspan=2, sticky="ew", padx=0, pady=0) + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=0) + + components.switch(frame, 0, 0, self.optimizer_ui_state, key, command=self.update_user_pref) + + self.muon_adam_button = components.button( + frame, 0, 1, "...", self.open_muon_adam_window, + tooltip="Configure the auxiliary AdamW_adv optimizer", + width=20, padx=5 ) + self.toggle_muon_adam_button() + elif type != 'bool': + components.entry(master, row, col + 1, self.optimizer_ui_state, key, + command=self.update_user_pref) + else: + components.switch(master, row, col + 1, self.optimizer_ui_state, key, + command=self.update_user_pref) + + def update_user_pref(self, *args): + update_optimizer_config(self.train_config) + self.toggle_muon_adam_button() + + def on_optimizer_change(self, *args): + optimizer_config = change_optimizer(self.train_config) + self.ui_state.get_var("optimizer").update(optimizer_config) + + self.clear_dynamic_ui(self.frame) + self.create_dynamic_ui(self.frame) + + def load_defaults(self, *args): + optimizer_config = load_optimizer_defaults(self.train_config) + self.ui_state.get_var("optimizer").update(optimizer_config) + + def on_window_close(self): + self.destroy() + + def toggle_muon_adam_button(self): + if self.muon_adam_button and self.muon_adam_button.winfo_exists(): + muon_with_adam = self.optimizer_ui_state.get_var("MuonWithAuxAdam").get() + self.muon_adam_button.configure(state="normal" if muon_with_adam else "disabled") + + def open_muon_adam_window(self): + current_optimizer = self.train_config.optimizer.optimizer + + adam_config = TrainOptimizerConfig.default_values() + current_state = self.train_config.optimizer.muon_adam_config + + if current_optimizer == Optimizer.MUON: + defaults = MUON_AUX_ADAM_DEFAULTS + else: + defaults = OPTIMIZER_DEFAULT_PARAMETERS[Optimizer.ADAMW_ADV] + + if current_state is None: + adam_config.from_dict(defaults) + if current_optimizer != Optimizer.MUON: + adam_config.optimizer = Optimizer.ADAMW_ADV + elif isinstance(current_state, dict): + adam_config.from_dict(current_state) + else: + # Should not happen if TrainConfig defines it as dict, but for safety + adam_config = current_state + + temp_adam_ui_state = UIState(self, adam_config) + window = MuonAdamWindow(self, self.train_config, temp_adam_ui_state, current_optimizer) + self.wait_window(window) + + self.train_config.optimizer.muon_adam_config = adam_config.to_dict() diff --git a/modules/ui/CtkProfilingWindowView.py b/modules/ui/CtkProfilingWindowView.py new file mode 100644 index 000000000..8d298abe3 --- /dev/null +++ b/modules/ui/CtkProfilingWindowView.py @@ -0,0 +1,57 @@ +import faulthandler + +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk +from scalene import scalene_profiler + + +class ProfilingWindow(ctk.CTkToplevel): + def __init__(self, parent, *args, **kwargs): + super().__init__(parent, *args, **kwargs) + self.parent = parent + + self.title("Profiling") + self.geometry("512x512") + self.resizable(True, True) + self.wait_visibility() + self.focus_set() + + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=0) + self.grid_rowconfigure(2, weight=1) + self.grid_columnconfigure(0, weight=1) + + components.button(self, 0, 0, "Dump stack", self._dump_stack) + self._profile_button = components.button( + self, 1, 0, "Start Profiling", self._start_profiler, + tooltip="Turns on/off Scalene profiling. Only works when OneTrainer is launched with Scalene!") + + # Bottom bar + self._bottom_bar = ctk.CTkFrame(master=self, corner_radius=0) + self._bottom_bar.grid(row=2, column=0, sticky="sew") + self._message_label = components.label(self._bottom_bar, 0, 0, "Inactive") + + self.protocol("WM_DELETE_WINDOW", self.withdraw) + self.withdraw() + self.after(200, lambda: set_window_icon(self)) + + def _dump_stack(self): + with open('stacks.txt', 'w') as f: + faulthandler.dump_traceback(f) + self._message_label.configure(text='Stack dumped to stacks.txt') + + def _end_profiler(self): + scalene_profiler.stop() + + self._message_label.configure(text='Inactive') + self._profile_button.configure(text='Start Profiling') + self._profile_button.configure(command=self._start_profiler) + + def _start_profiler(self): + scalene_profiler.start() + + self._message_label.configure(text='Profiling active...') + self._profile_button.configure(text='End Profiling') + self._profile_button.configure(command=self._end_profiler) diff --git a/modules/ui/CtkSampleFrameView.py b/modules/ui/CtkSampleFrameView.py new file mode 100644 index 000000000..297caac29 --- /dev/null +++ b/modules/ui/CtkSampleFrameView.py @@ -0,0 +1,134 @@ +from modules.util.config.SampleConfig import SampleConfig +from modules.util.enum.ModelType import ModelType +from modules.util.enum.NoiseScheduler import NoiseScheduler +from modules.util.ui import components +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class SampleFrame(ctk.CTkFrame): + def __init__( + self, + parent, + sample: SampleConfig, + ui_state: UIState, + model_type: ModelType, + include_prompt: bool = True, + include_settings: bool = True, + ): + ctk.CTkFrame.__init__(self, parent, fg_color="transparent") + + self.sample = sample + self.ui_state = ui_state + self.model_type = model_type + + is_flow_matching = model_type.is_flow_matching() + is_inpainting_model = model_type.has_conditioning_image_input() + is_video_model = model_type.is_video_model() + + if include_prompt and include_prompt: + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=1) + self.grid_columnconfigure(0, weight=1) + + if include_prompt: + top_frame = ctk.CTkFrame(self, fg_color="transparent") + top_frame.grid(row=0, column=0, padx=0, pady=0, sticky="nsew") + + top_frame.grid_columnconfigure(0, weight=0) + top_frame.grid_columnconfigure(1, weight=1) + + if include_settings: + bottom_frame = ctk.CTkFrame(self, fg_color="transparent") + bottom_frame.grid(row=1, column=0, padx=0, pady=0, sticky="nsew") + + bottom_frame.grid_columnconfigure(0, weight=0) + bottom_frame.grid_columnconfigure(1, weight=1) + bottom_frame.grid_columnconfigure(2, weight=0) + bottom_frame.grid_columnconfigure(3, weight=1) + + if include_prompt: + # prompt + components.label(top_frame, 0, 0, "prompt:") + components.entry(top_frame, 0, 1, self.ui_state, "prompt") + + # negative prompt + components.label(top_frame, 1, 0, "negative prompt:") + components.entry(top_frame, 1, 1, self.ui_state, "negative_prompt") + + if include_settings: + # width + components.label(bottom_frame, 0, 0, "width:") + components.entry(bottom_frame, 0, 1, self.ui_state, "width") + + # height + components.label(bottom_frame, 0, 2, "height:") + components.entry(bottom_frame, 0, 3, self.ui_state, "height") + + if is_video_model: + # frames + components.label(bottom_frame, 1, 0, "frames:", + tooltip="Number of frames to generate. Only used when generating videos.") + components.entry(bottom_frame, 1, 1, self.ui_state, "frames") + + # length + components.label(bottom_frame, 1, 2, "length:", + tooltip="Length in seconds of audio output.") + components.entry(bottom_frame, 1, 3, self.ui_state, "length") + + # seed + components.label(bottom_frame, 2, 0, "seed:") + components.entry(bottom_frame, 2, 1, self.ui_state, "seed") + + # random seed + components.label(bottom_frame, 2, 2, "random seed:") + components.switch(bottom_frame, 2, 3, self.ui_state, "random_seed") + + # cfg scale + components.label(bottom_frame, 3, 0, "cfg scale:") + components.entry(bottom_frame, 3, 1, self.ui_state, "cfg_scale") + + # sampler + if not is_flow_matching: + components.label(bottom_frame, 4, 2, "sampler:") + components.options_kv(bottom_frame, 4, 3, [ + ("DDIM", NoiseScheduler.DDIM), + ("Euler", NoiseScheduler.EULER), + ("Euler A", NoiseScheduler.EULER_A), + # ("DPM++", NoiseScheduler.DPMPP), # TODO: produces noisy samples + # ("DPM++ SDE", NoiseScheduler.DPMPP_SDE), # TODO: produces noisy samples + ("UniPC", NoiseScheduler.UNIPC), + ("Euler Karras", NoiseScheduler.EULER_KARRAS), + ("DPM++ Karras", NoiseScheduler.DPMPP_KARRAS), + ("DPM++ SDE Karras", NoiseScheduler.DPMPP_SDE_KARRAS), + ("UniPC Karras", NoiseScheduler.UNIPC_KARRAS) + ], self.ui_state, "noise_scheduler") + + # steps + components.label(bottom_frame, 4, 0, "steps:") + components.entry(bottom_frame, 4, 1, self.ui_state, "diffusion_steps") + + # inpainting + if is_inpainting_model: + components.label(bottom_frame, 5, 0, "inpainting:", + tooltip="Enables inpainting sampling. Only available when sampling from an inpainting model.") + components.switch(bottom_frame, 5, 1, self.ui_state, "sample_inpainting") + + # 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", + mode="file", + allow_model_files=False, + allow_image_files=True, + ) + + # 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", + mode="file", + allow_model_files=False, + allow_image_files=True, + ) diff --git a/modules/ui/CtkSampleParamsWindowView.py b/modules/ui/CtkSampleParamsWindowView.py new file mode 100644 index 000000000..2b0b3f3f1 --- /dev/null +++ b/modules/ui/CtkSampleParamsWindowView.py @@ -0,0 +1,39 @@ +from modules.ui.SampleFrame import SampleFrame +from modules.util.config.SampleConfig import SampleConfig +from modules.util.enum.ModelType import ModelType +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class SampleParamsWindow(ctk.CTkToplevel): + def __init__(self, parent, sample: SampleConfig, ui_state: UIState, model_type: ModelType | None = None, *args, **kwargs): + super().__init__(parent, *args, **kwargs) + + self.sample = sample + self.ui_state = ui_state + self.model_type = model_type + + self.title("Sample") + self.geometry("800x500") + self.resizable(True, True) + + self.grid_rowconfigure(0, weight=1) + self.grid_rowconfigure(1, weight=0) + self.grid_columnconfigure(0, weight=1) + + frame = SampleFrame(self, self.sample, self.ui_state, model_type=model_type) + frame.grid(row=0, column=0, padx=0, pady=0, sticky="nsew") + + components.button(self, 1, 0, "ok", self.__ok) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def __ok(self): + self.destroy() diff --git a/modules/ui/CtkSampleWindowView.py b/modules/ui/CtkSampleWindowView.py new file mode 100644 index 000000000..0f91ad2fa --- /dev/null +++ b/modules/ui/CtkSampleWindowView.py @@ -0,0 +1,227 @@ +import contextlib +import copy +import os +import tkinter as tk +import traceback + +from modules.model.BaseModel import BaseModel +from modules.modelSampler.BaseModelSampler import ( + BaseModelSampler, + ModelSamplerOutput, +) +from modules.ui.SampleFrame import SampleFrame +from modules.util import create +from modules.util.callbacks.TrainCallbacks import TrainCallbacks +from modules.util.commands.TrainCommands import TrainCommands +from modules.util.config.SampleConfig import SampleConfig +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.EMAMode import EMAMode +from modules.util.enum.FileType import FileType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.time_util import get_string_timestamp +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import torch + +import customtkinter as ctk +from PIL import Image + + +class SampleWindow(ctk.CTkToplevel): + def __init__( + self, + parent, + train_config: TrainConfig, + use_external_model: bool, + callbacks: TrainCallbacks | None = None, + commands: TrainCommands | None = None, + *args, **kwargs + ): + super().__init__(parent, *args, **kwargs) + + self.title("Sample") + self.geometry("1200x800") + self.resizable(True, True) + + if not use_external_model: + self.initial_train_config = TrainConfig.default_values().from_dict(train_config.to_dict()) + # remove some settings to speed up model loading for sampling + self.initial_train_config.optimizer.optimizer = None + self.initial_train_config.ema = EMAMode.OFF + else: + self.initial_train_config = None + + #TODO why is there a current_train_config and an initial_train_config? + #current_train_config doesn't seem to ever change + self.current_train_config = train_config + self.callbacks = callbacks + self.commands = commands + + # get model specific defaults + model_type = train_config.model_type + self.sample = SampleConfig.default_values(model_type) + self.ui_state = UIState(self, self.sample) + + if use_external_model: + self.callbacks.set_on_sample_custom(self.__update_preview) + self.callbacks.set_on_update_sample_custom_progress(self.__update_progress) + else: + self.model = None + self.model_sampler = None + + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=1) + self.grid_rowconfigure(2, weight=0) + self.grid_rowconfigure(3, weight=0) + self.grid_columnconfigure(0, weight=0) + self.grid_columnconfigure(1, weight=1) + + prompt_frame = SampleFrame(self, self.sample, self.ui_state, include_settings=False, model_type=model_type) + prompt_frame.grid(row=0, column=0, columnspan=2, padx=0, pady=0, sticky="nsew") + + settings_frame = SampleFrame(self, self.sample, self.ui_state, include_prompt=False, model_type=model_type) + settings_frame.grid(row=1, column=0, padx=0, pady=0, sticky="nsew") + + # image + self.image = ctk.CTkImage( + light_image=self.__dummy_image(), + size=(512, 512) + ) + + image_label = ctk.CTkLabel(master=self, text="", image=self.image, height=512, width=512) + image_label.grid(row=1, column=1, rowspan=3, sticky="nsew") + + self.progress = components.progress(self, 2, 0) + components.button(self, 3, 0, "sample", self.__sample) + + self.wait_visibility() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + def __load_model(self) -> BaseModel: + model_loader = create.create_model_loader( + model_type=self.initial_train_config.model_type, + training_method=self.initial_train_config.training_method, + ) + + model_setup = create.create_model_setup( + model_type=self.initial_train_config.model_type, + train_device=torch.device(self.initial_train_config.train_device), + temp_device=torch.device(self.initial_train_config.temp_device), + training_method=self.initial_train_config.training_method, + ) + + model_names = self.initial_train_config.model_names() + if self.initial_train_config.continue_last_backup: + last_backup_path = self.initial_train_config.get_last_backup_path() + + if last_backup_path: + if self.initial_train_config.training_method == TrainingMethod.LORA: + model_names.lora = last_backup_path + elif self.initial_train_config.training_method == TrainingMethod.EMBEDDING: + model_names.embedding.model_name = last_backup_path + else: # fine-tunes + model_names.base_model = last_backup_path + + print(f"Loading from backup '{last_backup_path}'...") + else: + print("No backup found, loading without backup...") + + if self.initial_train_config.quantization.cache_dir is None: + self.initial_train_config.quantization.cache_dir = self.initial_train_config.cache_dir + "/quantization" + os.makedirs(self.initial_train_config.quantization.cache_dir, exist_ok=True) + + model = model_loader.load( + model_type=self.initial_train_config.model_type, + model_names=model_names, + weight_dtypes=self.initial_train_config.weight_dtypes(), + quantization=self.initial_train_config.quantization, + ) + model.train_config = self.initial_train_config + + model_setup.setup_optimizations(model, self.initial_train_config) + model_setup.setup_train_device(model, self.initial_train_config) + model_setup.setup_model(model, self.initial_train_config) + model.to(torch.device(self.initial_train_config.temp_device)) + + return model + + def __create_sampler(self, model: BaseModel) -> BaseModelSampler: + return create.create_model_sampler( + train_device=torch.device(self.initial_train_config.train_device), + temp_device=torch.device(self.initial_train_config.temp_device), + model=model, + model_type=self.initial_train_config.model_type, + training_method=self.initial_train_config.training_method, + ) + + def __update_preview(self, sampler_output: ModelSamplerOutput): + if sampler_output.file_type == FileType.IMAGE: + image = sampler_output.data + self.image.configure( + light_image=image, + size=(image.width, image.height), + ) + + def __update_progress(self, progress: int, max_progress: int): + self.progress.set(progress / max_progress) + self.update() + + def __dummy_image(self) -> Image: + return Image.new(mode="RGB", size=(512, 512), color=(0, 0, 0)) + + def __sample(self): + sample = copy.copy(self.sample) + + if self.commands: + self.commands.sample_custom(sample) + else: + if self.model is None: + # lazy initialization + self.model = self.__load_model() + self.model_sampler = self.__create_sampler(self.model) + + sample.from_train_config(self.current_train_config) + + sample_dir = os.path.join( + self.initial_train_config.workspace_dir, + "samples", + "custom", + ) + + progress = self.model.train_progress + sample_path = os.path.join( + sample_dir, + f"{get_string_timestamp()}-training-sample-{progress.filename_string()}" + ) + + self.model.eval() + + self.model_sampler.sample( + sample_config=sample, + destination=sample_path, + image_format=self.current_train_config.sample_image_format, + video_format=self.current_train_config.sample_video_format, + audio_format=self.current_train_config.sample_audio_format, + on_sample=self.__update_preview, + on_update_progress=self.__update_progress, + ) + + def destroy(self): + try: + if hasattr(self, "_icon_image_ref"): + del self._icon_image_ref + + # Remove any pending after callbacks + for after_id in self.tk.call('after', 'info'): + with contextlib.suppress(tk.TclError, RuntimeError): + self.after_cancel(after_id) + + super().destroy() + except (tk.TclError, RuntimeError) as e: + print(f"Error destroying window: {e}") + except Exception as e: + print(f"Unexpected error destroying window: {e}") + traceback.print_exc() diff --git a/modules/ui/CtkSamplingTabView.py b/modules/ui/CtkSamplingTabView.py new file mode 100644 index 000000000..5a3c44f08 --- /dev/null +++ b/modules/ui/CtkSamplingTabView.py @@ -0,0 +1,124 @@ +from modules.ui.ConfigList import ConfigList +from modules.ui.SampleParamsWindow import SampleParamsWindow +from modules.util.config.SampleConfig import SampleConfig +from modules.util.config.TrainConfig import TrainConfig +from modules.util.ui import components +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class SamplingTab(ConfigList): + + def __init__(self, master, train_config: TrainConfig, ui_state: UIState): + super().__init__( + master, + train_config, + ui_state, + from_external_file=True, + attr_name="sample_definition_file_name", + config_dir="training_samples", + default_config_name="samples.json", + add_button_text="Add Sample", + add_button_tooltip="Add a new sample configuration.", + is_full_width=True, + show_toggle_button=True + ) + + def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): + return SampleWidget(master, element, i, open_command, remove_command, clone_command, save_command) + + def create_new_element(self) -> dict: + return SampleConfig.default_values(self.train_config.model_type) + + def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: + return SampleParamsWindow(self.master, self.current_config[i], ui_state, model_type=self.train_config.model_type) + + +class SampleWidget(ctk.CTkFrame): + def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): + super().__init__( + master=master, corner_radius=10, bg_color="transparent" + ) + + self.element = element + self.ui_state = UIState(self, element) + self.i = i + self.save_command = save_command + + self.grid_columnconfigure(10, weight=1) + + # close button + close_button = ctk.CTkButton( + master=self, + width=20, + height=20, + text="X", + corner_radius=2, + fg_color="#C00000", + command=lambda: remove_command(self.i), + ) + close_button.grid(row=0, column=0) + + # clone button + clone_button = ctk.CTkButton( + master=self, + width=20, + height=20, + text="+", + corner_radius=2, + fg_color="#00C000", + command=lambda: clone_command(self.i), + ) + clone_button.grid(row=0, column=1, padx=5) + + # enabled + self.enabled_switch = components.switch(self, 0, 2, self.ui_state, "enabled", self.__switch_enabled) + self.enabled_switch.configure(width=40) + + # width + components.label(self, 0, 3, "width:") + self.width_entry = components.entry(self, 0, 4, self.ui_state, "width") + self.width_entry.bind('', lambda _: save_command()) + self.width_entry.configure(width=50) + + # height + components.label(self, 0, 5, "height:") + self.height_entry = components.entry(self, 0, 6, self.ui_state, "height") + self.height_entry.bind('', lambda _: save_command()) + self.height_entry.configure(width=50) + + # seed + components.label(self, 0, 7, "seed:") + self.seed_entry = components.entry(self, 0, 8, self.ui_state, "seed") + self.seed_entry.bind('', lambda _: save_command()) + self.seed_entry.configure(width=80) + + # prompt + components.label(self, 0, 9, "prompt:") + self.prompt_entry = components.entry(self, 0, 10, self.ui_state, "prompt") + self.prompt_entry.bind('', lambda _: save_command()) + + # button + self.button = components.icon_button(self, 0, 11, "...", lambda: open_command(self.i, self.ui_state)) + self.button.configure(width=40) + + self.__set_enabled() + + def __switch_enabled(self): + self.save_command() + self.__set_enabled() + + def __set_enabled(self): + enabled = self.element.enabled + self.width_entry.configure(state="normal" if enabled else "disabled") + self.height_entry.configure(state="normal" if enabled else "disabled") + self.prompt_entry.configure(state="normal" if enabled else "disabled") + self.seed_entry.configure(state="normal" if enabled else "disabled") + self.button.configure(state="normal" if enabled else "disabled") + + def configure_element(self): + pass + + def place_in_list(self): + self.grid(row=self.i, column=0, pady=5, padx=5, sticky="new") diff --git a/modules/ui/CtkSchedulerParamsWindowView.py b/modules/ui/CtkSchedulerParamsWindowView.py new file mode 100644 index 000000000..f96ed4876 --- /dev/null +++ b/modules/ui/CtkSchedulerParamsWindowView.py @@ -0,0 +1,119 @@ +from modules.ui.ConfigList import ConfigList +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.LearningRateScheduler import LearningRateScheduler +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class KvParams(ConfigList): + def __init__(self, master, train_config: TrainConfig, ui_state: UIState): + super().__init__( + master, + train_config, + ui_state, + attr_name="scheduler_params", + from_external_file=False, + add_button_text="add parameter", + is_full_width=True + ) + + def refresh_ui(self): + self._create_element_list() + + def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): + return KvWidget(master, element, i, open_command, remove_command, clone_command, save_command) + + def create_new_element(self) -> dict[str, str]: + return {"key": "", "value": ""} + + def open_element_window(self, i, ui_state): + pass + + +class KvWidget(ctk.CTkFrame): + def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): + super().__init__(master=master, bg_color="transparent") + self.element = element + self.ui_state = UIState(self, element) + self.i = i + self.save_command = save_command + + self.grid_columnconfigure(0, weight=0) + self.grid_columnconfigure(1, weight=1, uniform=1) + self.grid_columnconfigure(2, weight=1, uniform=1) + + close_button = ctk.CTkButton( + master=self, + width=20, + height=20, + text="X", + corner_radius=2, + fg_color="#C00000", + command=lambda: remove_command(self.i)) + close_button.grid(row=0, column=0) + + # Key + tooltip_key = "Key name for an argument in your scheduler" + self.key = components.entry(self, 0, 1, self.ui_state, "key", + tooltip=tooltip_key, wide_tooltip=True) + self.key.bind("", lambda _: save_command()) + self.key.configure(width=50) + + # Value + tooltip_val = "Value for an argument in your scheduler. Some special values can be used, wrapped in percent signs: LR, EPOCHS, STEPS_PER_EPOCH, TOTAL_STEPS, SCHEDULER_STEPS. Note that OneTrainer calls step() after every individual learning step, not every epoch, so what Torch calls 'epoch' you should treat as 'step'." + self.value = components.entry(self, 0, 2, self.ui_state, "value", + tooltip=tooltip_val, wide_tooltip=True) + self.value.bind("", lambda _: save_command()) + self.value.configure(width=50) + + def place_in_list(self): + self.grid(row=self.i, column=0, padx=5, pady=5, sticky="new") + + +class SchedulerParamsWindow(ctk.CTkToplevel): + def __init__(self, parent, train_config: TrainConfig, ui_state, *args, **kwargs): + super().__init__(parent, *args, **kwargs) + + self.parent = parent + self.train_config = train_config + self.ui_state = ui_state + + self.title("Learning Rate Scheduler Settings") + self.geometry("800x400") + self.resizable(True, True) + + self.grid_rowconfigure(0, weight=1) + self.grid_rowconfigure(1, weight=0) + self.grid_columnconfigure(0, weight=1) + + self.frame = ctk.CTkFrame(self) + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + self.frame.grid_columnconfigure(0, weight=0) + self.frame.grid_columnconfigure(1, weight=1) + + self.expand_frame = ctk.CTkFrame(self.frame, bg_color="transparent") + self.expand_frame.grid(row=1, column=0, columnspan=2, sticky="nsew") + + components.button(self, 1, 0, "ok", command=self.on_window_close) + self.main_frame(self.frame) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def main_frame(self, master): + if self.train_config.learning_rate_scheduler is LearningRateScheduler.CUSTOM: + components.label(master, 0, 0, "Class Name", + tooltip="Python class module and name for the custom scheduler class, in the form of ..") + components.entry(master, 0, 1, self.ui_state, "custom_learning_rate_scheduler") + + # Any additional parameters, in key-value form. + self.params = KvParams(self.expand_frame, self.train_config, self.ui_state) + + def on_window_close(self): + self.destroy() diff --git a/modules/ui/CtkTimestepDistributionWindowView.py b/modules/ui/CtkTimestepDistributionWindowView.py new file mode 100644 index 000000000..21e41ce3e --- /dev/null +++ b/modules/ui/CtkTimestepDistributionWindowView.py @@ -0,0 +1,186 @@ + +from modules.modelSetup.mixin.ModelSetupNoiseMixin import ( + ModelSetupNoiseMixin, +) +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.TimestepDistribution import TimestepDistribution +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import torch +from torch import Tensor + +import customtkinter as ctk +from customtkinter import AppearanceModeTracker, ThemeManager +from matplotlib import pyplot as plt +from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg + + +class TimestepGenerator(ModelSetupNoiseMixin): + + def __init__( + self, + timestep_distribution: TimestepDistribution, + min_noising_strength: float, + max_noising_strength: float, + noising_weight: float, + noising_bias: float, + timestep_shift: float, + ): + super().__init__() + + self.timestep_distribution = timestep_distribution + self.min_noising_strength = min_noising_strength + self.max_noising_strength = max_noising_strength + self.noising_weight = noising_weight + self.noising_bias = noising_bias + self.timestep_shift = timestep_shift + + def generate(self) -> Tensor: + generator = torch.Generator() + generator.seed() + + config = TrainConfig.default_values() + config.timestep_distribution = self.timestep_distribution + config.min_noising_strength = self.min_noising_strength + config.max_noising_strength = self.max_noising_strength + config.noising_weight = self.noising_weight + config.noising_bias = self.noising_bias + config.timestep_shift = self.timestep_shift + + + return self._get_timestep_discrete( + num_train_timesteps=1000, + deterministic=False, + generator=generator, + batch_size=1000000, + config=config, + ) + + +class TimestepDistributionWindow(ctk.CTkToplevel): + def __init__( + self, + parent, + config: TrainConfig, + ui_state: UIState, + *args, **kwargs, + ): + super().__init__(parent, *args, **kwargs) + + self.title("Timestep Distribution") + self.geometry("900x600") + self.resizable(True, True) + + self.config = config + self.ui_state = ui_state + self.image_preview_file_index = 0 + self.ax = None + self.canvas = None + + self.grid_rowconfigure(0, weight=1) + self.grid_columnconfigure(0, weight=1) + + frame = self.__content_frame(self) + frame.grid(row=0, column=0, sticky='nsew') + components.button(self, 1, 0, "ok", self.__ok) + + self.wait_visibility() + self.after(200, lambda: set_window_icon(self)) + self.grab_set() + self.focus_set() + + def __content_frame(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=0) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + frame.grid_rowconfigure(7, weight=1) + + # timestep distribution + components.label(frame, 0, 0, "Timestep Distribution", + tooltip="Selects the function to sample timesteps during training", + wide_tooltip=True) + components.options(frame, 0, 1, [str(x) for x in list(TimestepDistribution)], self.ui_state, + "timestep_distribution") + + # min noising strength + components.label(frame, 1, 0, "Min Noising Strength", + tooltip="Specifies the minimum noising strength used during training. This can help to improve composition, but prevents finer details from being trained") + components.entry(frame, 1, 1, self.ui_state, "min_noising_strength") + + # max noising strength + components.label(frame, 2, 0, "Max Noising Strength", + tooltip="Specifies the maximum noising strength used during training. This can be useful to reduce overfitting, but also reduces the impact of training samples on the overall image composition") + components.entry(frame, 2, 1, self.ui_state, "max_noising_strength") + + # noising weight + components.label(frame, 3, 0, "Noising Weight", + tooltip="Controls the weight parameter of the timestep distribution function. Use the preview to see more details.") + components.entry(frame, 3, 1, self.ui_state, "noising_weight") + + # noising bias + components.label(frame, 4, 0, "Noising Bias", + tooltip="Controls the bias parameter of the timestep distribution function. Use the preview to see more details.") + components.entry(frame, 4, 1, self.ui_state, "noising_bias") + + # timestep shift + components.label(frame, 5, 0, "Timestep Shift", + tooltip="Shift the timestep distribution. Use the preview to see more details.") + components.entry(frame, 5, 1, self.ui_state, "timestep_shift") + + # dynamic timestep shifting + components.label(frame, 6, 0, "Dynamic Timestep Shifting", + tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Dynamic Timestep Shifting is not shown in the preview. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) + components.switch(frame, 6, 1, self.ui_state, "dynamic_timestep_shifting") + + + # plot + appearance_mode = AppearanceModeTracker.get_mode() + background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) + text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) + background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" + text_color = f"#{int(text_color[0]/256):x}{int(text_color[1]/256):x}{int(text_color[2]/256):x}" + + fig, ax = plt.subplots() + self.ax = ax + self.canvas = FigureCanvasTkAgg(fig, master=frame) + self.canvas.get_tk_widget().grid(row=0, column=3, rowspan=8) + + fig.set_facecolor(background_color) + ax.set_facecolor(background_color) + ax.spines['bottom'].set_color(text_color) + ax.spines['left'].set_color(text_color) + ax.spines['top'].set_color(text_color) + ax.spines['right'].set_color(text_color) + ax.tick_params(axis='x', colors=text_color, which="both") + ax.tick_params(axis='y', colors=text_color, which="both") + ax.xaxis.label.set_color(text_color) + ax.yaxis.label.set_color(text_color) + + self.__update_preview() + + # update button + components.button(frame, 8, 3, "Update Preview", command=self.__update_preview) + + frame.pack(fill="both", expand=1) + return frame + + def __update_preview(self): + generator = TimestepGenerator( + timestep_distribution=self.config.timestep_distribution, + min_noising_strength=self.config.min_noising_strength, + max_noising_strength=self.config.max_noising_strength, + noising_weight=self.config.noising_weight, + noising_bias=self.config.noising_bias, + timestep_shift=self.config.timestep_shift, + ) + + self.ax.cla() + self.ax.hist(generator.generate(), bins=1000, range=(0, 999)) + self.canvas.draw() + + def __ok(self): + self.destroy() diff --git a/modules/ui/CtkTopBarView.py b/modules/ui/CtkTopBarView.py new file mode 100644 index 000000000..820fdb71a --- /dev/null +++ b/modules/ui/CtkTopBarView.py @@ -0,0 +1,260 @@ +import json +import os +import traceback +import webbrowser +from collections.abc import Callable +from contextlib import suppress + +from modules.util import path_util +from modules.util.config.SecretsConfig import SecretsConfig +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ModelType import ModelType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.optimizer_util import change_optimizer +from modules.util.path_util import write_json_atomic +from modules.util.ui import components, dialogs +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class TopBar: + def __init__( + self, + master, + train_config: TrainConfig, + ui_state: UIState, + change_model_type_callback: Callable[[ModelType], None], + change_training_method_callback: Callable[[TrainingMethod], None], + load_preset_callback: Callable[[], None], + ): + self.master = master + self.train_config = train_config + self.ui_state = ui_state + self.change_model_type_callback = change_model_type_callback + self.change_training_method_callback = change_training_method_callback + self.load_preset_callback = load_preset_callback + + self.dir = "training_presets" + + self.config_ui_data = { + "config_name": path_util.canonical_join(self.dir, "#.json") + } + 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.current_config = [] + + self.frame = ctk.CTkFrame(master=master, corner_radius=0) + self.frame.grid(row=0, column=0, sticky="nsew") + + self.training_method = None + + # title + components.app_title(self.frame, 0, 0) + + # dropdown + self.configs_dropdown = None + self.__create_configs_dropdown() + + # remove button + # TODO + # components.icon_button(self.frame, 0, 2, "-", self.__remove_config) + + # Wiki button + components.button(self.frame, 0, 4, "Wiki", self.open_wiki, width=50) + + # save button + components.button(self.frame, 0, 3, "Save config", self.__save_config, + tooltip="Save the current configuration in a custom preset", width=90) + + # padding + self.frame.grid_columnconfigure(5, weight=1) + + # model type + components.options_kv( + master=self.frame, + row=0, + column=6, + values=[ #TODO simplify + ("SD1.5", ModelType.STABLE_DIFFUSION_15), + ("SD1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), + ("SD2.0", ModelType.STABLE_DIFFUSION_20), + ("SD2.0 Inpainting", ModelType.STABLE_DIFFUSION_20_INPAINTING), + ("SD2.1", ModelType.STABLE_DIFFUSION_21), + ("SD3", ModelType.STABLE_DIFFUSION_3), + ("SD3.5", ModelType.STABLE_DIFFUSION_35), + ("SDXL", ModelType.STABLE_DIFFUSION_XL_10_BASE), + ("SDXL Inpainting", ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING), + ("Wuerstchen v2", ModelType.WUERSTCHEN_2), + ("Stable Cascade", ModelType.STABLE_CASCADE_1), + ("PixArt Alpha", ModelType.PIXART_ALPHA), + ("PixArt Sigma", ModelType.PIXART_SIGMA), + ("Flux Dev.1", ModelType.FLUX_DEV_1), + ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), + ("Flux 2 [Dev, Klein]", ModelType.FLUX_2), + ("Sana", ModelType.SANA), + ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), + ("HiDream Full", ModelType.HI_DREAM_FULL), + ("Chroma1", ModelType.CHROMA_1), + ("QwenImage", ModelType.QWEN), + ("Z-Image", ModelType.Z_IMAGE), + ("Ernie Image", ModelType.ERNIE), + ], + ui_state=self.ui_state, + var_name="model_type", + command=self.__change_model_type, + ) + + def __create_training_method(self): + if self.training_method: + self.training_method.destroy() + + values = [] + #TODO simplify + if self.train_config.model_type.is_stable_diffusion(): + values = [ + ("Fine Tune", TrainingMethod.FINE_TUNE), + ("LoRA", TrainingMethod.LORA), + ("Embedding", TrainingMethod.EMBEDDING), + ("Fine Tune VAE", TrainingMethod.FINE_TUNE_VAE), + ] + elif self.train_config.model_type.is_stable_diffusion_3() \ + or self.train_config.model_type.is_stable_diffusion_xl() \ + or self.train_config.model_type.is_wuerstchen() \ + or self.train_config.model_type.is_pixart() \ + or self.train_config.model_type.is_flux_1() \ + or self.train_config.model_type.is_sana() \ + or self.train_config.model_type.is_hunyuan_video() \ + or self.train_config.model_type.is_hi_dream() \ + or self.train_config.model_type.is_chroma(): + values = [ + ("Fine Tune", TrainingMethod.FINE_TUNE), + ("LoRA", TrainingMethod.LORA), + ("Embedding", TrainingMethod.EMBEDDING), + ] + elif self.train_config.model_type.is_qwen() \ + or self.train_config.model_type.is_z_image() \ + or self.train_config.model_type.is_flux_2() \ + or self.train_config.model_type.is_ernie(): + values = [ + ("Fine Tune", TrainingMethod.FINE_TUNE), + ("LoRA", TrainingMethod.LORA), + ] + + # training method + self.training_method = components.options_kv( + master=self.frame, + row=0, + column=7, + values=values, + ui_state=self.ui_state, + var_name="training_method", + command=self.change_training_method_callback, + ) + + def __change_model_type(self, model_type: ModelType): + self.change_model_type_callback(model_type) + self.__create_training_method() + + def __create_configs_dropdown(self): + if self.configs_dropdown is not None: + self.configs_dropdown.grid_forget() + + self.configs_dropdown = components.options_kv( + self.frame, 0, 1, self.configs, self.config_ui_state, "config_name", self.__load_current_config + ) + + def __load_available_config_names(self): + if os.path.isdir(self.dir): + for path in os.listdir(self.dir): + if path != "#.json": + path = path_util.canonical_join(self.dir, path) + if path.endswith(".json") and os.path.isfile(path): + name = os.path.basename(path) + name = os.path.splitext(name)[0] + self.configs.append((name, path)) + self.configs.sort() + + def __save_to_file(self, name) -> str: + name = path_util.safe_filename(name) + path = path_util.canonical_join("training_presets", f"{name}.json") + + write_json_atomic(path, self.train_config.to_settings_dict(secrets=False)) + + return path + + def __save_secrets(self, path) -> str: + write_json_atomic(path, self.train_config.secrets.to_dict()) + return path + + def open_wiki(self): + webbrowser.open("https://github.com/Nerogar/OneTrainer/wiki", new=0, autoraise=False) + + def __save_new_config(self, name): + path = self.__save_to_file(name) + + is_new_config = name not in [x[0] for x in self.configs] + + if is_new_config: + self.configs.append((name, path)) + self.configs.sort() + + if self.config_ui_data["config_name"] != path_util.canonical_join(self.dir, f"{name}.json"): + self.config_ui_state.get_var("config_name").set(path_util.canonical_join(self.dir, f"{name}.json")) + + if is_new_config: + self.__create_configs_dropdown() + + def __save_config(self): + default_value = self.configs_dropdown.get() + while default_value.startswith('#'): + default_value = default_value[1:] + + dialogs.StringInputDialog( + parent=self.master, + title="name", + question="Config Name", + callback=self.__save_new_config, + default_value=default_value, + validate_callback=lambda x: not x.startswith("#") + ) + + def __load_current_config(self, filename): + try: + basename = os.path.basename(filename) + is_built_in_preset = basename.startswith("#") and basename != "#.json" + + with open(filename, "r") as f: + loaded_dict = json.load(f) + default_config = TrainConfig.default_values() + if is_built_in_preset: + # always assume built-in configs are saved in the most recent version + loaded_dict["__version"] = default_config.config_version + loaded_config = default_config.from_dict(loaded_dict).to_unpacked_config() + + with suppress(FileNotFoundError), open("secrets.json", "r") as f: + secrets_dict=json.load(f) + loaded_config.secrets = SecretsConfig.default_values().from_dict(secrets_dict) + + self.train_config.from_dict(loaded_config.to_dict()) + self.ui_state.update(loaded_config) + + optimizer_config = change_optimizer(self.train_config) + self.ui_state.get_var("optimizer").update(optimizer_config) + + self.load_preset_callback() + except FileNotFoundError: + pass + except Exception: + print(traceback.format_exc()) + + def __remove_config(self): + # TODO + pass + + def save_default(self): + self.__save_to_file("#") + self.__save_secrets("secrets.json") diff --git a/modules/ui/CtkTrainUIView.py b/modules/ui/CtkTrainUIView.py new file mode 100644 index 000000000..ba90d2e64 --- /dev/null +++ b/modules/ui/CtkTrainUIView.py @@ -0,0 +1,889 @@ +import ctypes +import datetime +import json +import os +import platform +import subprocess +import sys +import threading +import time +import traceback +import webbrowser +from collections.abc import Callable +from contextlib import suppress +from pathlib import Path +from tkinter import filedialog, messagebox + +import scripts.generate_debug_report +from modules.ui.AdditionalEmbeddingsTab import AdditionalEmbeddingsTab +from modules.ui.CaptionUI import CaptionUI +from modules.ui.CloudTab import CloudTab +from modules.ui.ConceptTab import ConceptTab +from modules.ui.ConvertModelUI import ConvertModelUI +from modules.ui.LoraTab import LoraTab +from modules.ui.ModelTab import ModelTab +from modules.ui.ProfilingWindow import ProfilingWindow +from modules.ui.SampleWindow import SampleWindow +from modules.ui.SamplingTab import SamplingTab +from modules.ui.TopBar import TopBar +from modules.ui.TrainingTab import TrainingTab +from modules.ui.VideoToolUI import VideoToolUI +from modules.util import create +from modules.util.callbacks.TrainCallbacks import TrainCallbacks +from modules.util.commands.TrainCommands import TrainCommands +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.DataType import DataType +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.PathIOType import PathIOType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.torch_util import torch_gc +from modules.util.TrainProgress import TrainProgress +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState +from modules.util.ui.validation import flush_and_validate_all + +import torch + +import customtkinter as ctk +from customtkinter import AppearanceModeTracker + +# 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": + with suppress(Exception): + # 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): + set_step_progress: Callable[[int, int], None] + set_epoch_progress: Callable[[int, int], None] + + status_label: ctk.CTkLabel | None + training_button: ctk.CTkButton | None + training_callbacks: TrainCallbacks | None + training_commands: TrainCommands | None + + _TRAIN_BUTTON_STYLES = { + "idle": { + "text": "Start Training", + "state": "normal", + "fg_color": "#198754", + "hover_color": "#146c43", + "text_color": "white", + "text_color_disabled": "white", + }, + "running": { + "text": "Stop Training", + "state": "normal", + "fg_color": "#dc3545", + "hover_color": "#bb2d3b", + "text_color": "white", + }, + "stopping": { + "text": "Stopping...", + "state": "disabled", + "fg_color": "#dc3545", + "hover_color": "#dc3545", + "text_color": "white", + "text_color_disabled": "white", + }, + } + + def __init__(self): + super().__init__() + + self.title("OneTrainer") + self.geometry("1100x740") + + self.after(100, lambda: self._set_icon()) + + # more efficient version of ctk.set_appearance_mode("System"), which retrieves the system theme on each main loop iteration + ctk.set_appearance_mode("Light" if AppearanceModeTracker.detect_appearance_mode() == 0 else "Dark") + ctk.set_default_color_theme("blue") + + self.train_config = TrainConfig.default_values() + self.ui_state = UIState(self, self.train_config) + + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=1) + self.grid_rowconfigure(2, weight=0) + self.grid_columnconfigure(0, weight=1) + + self.status_label = None + self.eta_label = None + self.training_button = None + self.export_button = None + self.tabview = None + + self.model_tab = None + self.training_tab = None + self.lora_tab = None + self.cloud_tab = None + self.additional_embeddings_tab = None + + self.top_bar_component = self.top_bar(self) + self.content_frame(self) + self.bottom_bar(self) + + 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.workspace_dir_trace_id = self.ui_state.add_var_trace("workspace_dir", self._on_workspace_dir_change_trace) + + # Persistent profiling window. + self.profiling_window = ProfilingWindow(self) + + self.protocol("WM_DELETE_WINDOW", self.__close) + + def __close(self): + self.top_bar_component.save_default() + self._stop_always_on_tensorboard() + if hasattr(self, 'workspace_dir_trace_id'): + self.ui_state.remove_var_trace("workspace_dir", self.workspace_dir_trace_id) + self.quit() + + def top_bar(self, master): + return TopBar( + master, + self.train_config, + self.ui_state, + self.change_model_type, + self.change_training_method, + self.load_preset, + ) + + def _set_icon(self): + """Set the window icon safely after window is ready""" + set_window_icon(self) + + def bottom_bar(self, master): + frame = ctk.CTkFrame(master=master, corner_radius=0) + frame.grid(row=2, column=0, sticky="nsew") + + self.set_step_progress, self.set_epoch_progress = components.double_progress(frame, 0, 0, "step", "epoch") + + # status + ETA container + self.status_frame = ctk.CTkFrame(frame, corner_radius=0, fg_color="transparent") + self.status_frame.grid(row=0, column=1, sticky="w") + self.status_frame.grid_rowconfigure(0, weight=0) + self.status_frame.grid_rowconfigure(1, weight=0) + self.status_frame.grid_columnconfigure(0, weight=1) + + self.status_label = components.label(self.status_frame, 0, 0, "", pad=0, + tooltip="Current status of the training run") + self.eta_label = components.label(self.status_frame, 1, 0, "", pad=0) + + # padding + frame.grid_columnconfigure(2, weight=1) + + + # export button + self.export_button = components.button(frame, 0, 3, "Export", self.export_training, + width=60, padx=5, pady=(15, 0), + tooltip="Export the current configuration as a script to run without a UI") + + # debug button + components.button(frame, 0, 4, "Debug", self.generate_debug_package, + width=60, padx=(5, 25), pady=(15, 0), + tooltip="Generate a zip file with config.json, debug_report.log and settings diff, use this to report bugs or issues") + + # tensorboard button + components.button(frame, 0, 5, "Tensorboard", self.open_tensorboard, + width=100, padx=(0, 5), pady=(15, 0)) + + # training button + self.training_button = components.button(frame, 0, 6, "Start Training", self.start_training, + padx=(5, 20), pady=(15, 0)) + self._set_training_button_style("idle") # centralized styling + + return frame + + def content_frame(self, master): + frame = ctk.CTkFrame(master=master, corner_radius=0) + frame.grid(row=1, column=0, sticky="nsew") + + frame.grid_rowconfigure(0, weight=1) + frame.grid_columnconfigure(0, weight=1) + + self.tabview = ctk.CTkTabview(frame) + self.tabview.grid(row=0, column=0, sticky="nsew") + + 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) + + return frame + + def create_general_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # workspace dir + components.label(frame, 0, 0, "Workspace Directory", + tooltip="The directory where all files of this training run are saved") + components.path_entry(frame, 0, 1, self.ui_state, "workspace_dir", mode="dir", command=self._on_workspace_dir_change) + + # cache dir + components.label(frame, 0, 2, "Cache Directory", + tooltip="The directory where cached data is saved") + components.path_entry(frame, 0, 3, self.ui_state, "cache_dir", mode="dir") + + # continue from previous backup + components.label(frame, 2, 0, "Continue from last backup", + tooltip="Automatically continues training from the last backup saved in /backup") + components.switch(frame, 2, 1, self.ui_state, "continue_last_backup") + + # only cache + components.label(frame, 2, 2, "Only Cache", + tooltip="Only populate the cache, without any training") + components.switch(frame, 2, 3, self.ui_state, "only_cache") + + # TODO: In Phase 4 rework the general tab. + # prevent overwrites + components.label(frame, 3, 0, "Prevent Overwrites", + tooltip="When enabled, output paths that already exist on disk will be flagged as invalid to avoid accidental overwrites") + components.switch(frame, 3, 1, self.ui_state, "prevent_overwrites") + + # debug + components.label(frame, 4, 0, "Debug mode", + tooltip="Save debug information during the training into the debug directory") + components.switch(frame, 4, 1, self.ui_state, "debug_mode") + + components.label(frame, 4, 2, "Debug Directory", + tooltip="The directory where debug data is saved") + components.path_entry(frame, 4, 3, self.ui_state, "debug_dir", mode="dir", io_type=PathIOType.OUTPUT) + + # 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", + tooltip="Port to use for Tensorboard link") + components.entry(frame, 7, 3, self.ui_state, "tensorboard_port") + + + # validation + components.label(frame, 8, 0, "Validation", + tooltip="Enable validation steps and add new graph in tensorboard") + components.switch(frame, 8, 1, self.ui_state, "validation") + + components.label(frame, 8, 2, "Validate after", + tooltip="The interval used when validate training") + components.time_entry(frame, 8, 3, self.ui_state, "validate_after", "validate_after_unit") + + # device + components.label(frame, 10, 0, "Dataloader Threads", + tooltip="Number of threads used for the data loader. Increase if your GPU has room during caching, decrease if it's going out of memory during caching.") + components.entry(frame, 10, 1, self.ui_state, "dataloader_threads", required=True) + + components.label(frame, 11, 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", required=True) + + components.label(frame, 12, 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", + 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.label(frame, 13, 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, + "gradient_reduce_precision") + + components.label(frame, 13, 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.label(frame, 14, 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)", + 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.label(frame, 15, 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") + + frame.pack(fill="both", expand=1) + return frame + + def create_model_tab(self, master): + return ModelTab(master, self.train_config, self.ui_state) + + def create_data_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + + # aspect ratio bucketing + components.label(frame, 0, 0, "Aspect Ratio Bucketing", + tooltip="Aspect ratio bucketing enables training on images with different aspect ratios") + components.switch(frame, 0, 1, self.ui_state, "aspect_ratio_bucketing") + + # latent caching + components.label(frame, 1, 0, "Latent Caching", + tooltip="Caching of intermediate training data that can be re-used between epochs") + components.switch(frame, 1, 1, self.ui_state, "latent_caching") + + # clear cache before training + components.label(frame, 2, 0, "Clear cache before training", + tooltip="Clears the cache directory before starting to train. Only disable this if you want to continue using the same cached data. Disabling this can lead to errors, if other settings are changed during a restart") + components.switch(frame, 2, 1, self.ui_state, "clear_cache_before_training") + + frame.pack(fill="both", expand=1) + return frame + + def create_concepts_tab(self, master): + return ConceptTab(master, self.train_config, self.ui_state) + + def create_training_tab(self, master) -> TrainingTab: + return TrainingTab(master, self.train_config, self.ui_state) + + def create_cloud_tab(self, master) -> CloudTab: + return CloudTab(master, self.train_config, self.ui_state,parent=self) + + def create_sampling_tab(self, master): + master.grid_rowconfigure(0, weight=0) + master.grid_rowconfigure(1, weight=1) + master.grid_columnconfigure(0, weight=1) + + # sample after + top_frame = ctk.CTkFrame(master=master, corner_radius=0) + top_frame.grid(row=0, column=0, sticky="nsew") + sub_frame = ctk.CTkFrame(master=top_frame, corner_radius=0, fg_color="transparent") + sub_frame.grid(row=1, column=0, sticky="nsew", columnspan=6) + + components.label(top_frame, 0, 0, "Sample After", + tooltip="The interval used when automatically sampling from the model during training") + components.time_entry(top_frame, 0, 1, self.ui_state, "sample_after", "sample_after_unit") + + components.label(top_frame, 0, 2, "Skip First", + tooltip="Start sampling automatically after this interval has elapsed.") + components.entry(top_frame, 0, 3, self.ui_state, "sample_skip_first", width=50, sticky="nw") + + components.label(top_frame, 0, 4, "Format", + tooltip="File Format used when saving samples") + components.options_kv(top_frame, 0, 5, [ + ("PNG", ImageFormat.PNG), + ("JPG", ImageFormat.JPG), + ], self.ui_state, "sample_image_format") + + components.button(top_frame, 0, 6, "sample now", self.sample_now) + + components.button(top_frame, 0, 7, "manual sample", self.open_manual_sample_window ) + + components.label(sub_frame, 0, 0, "Non-EMA Sampling", + tooltip="Whether to include non-ema sampling when using ema.") + components.switch(sub_frame, 0, 1, self.ui_state, "non_ema_sampling") + + components.label(sub_frame, 0, 2, "Samples to Tensorboard", + tooltip="Whether to include sample images in the Tensorboard output.") + components.switch(sub_frame, 0, 3, self.ui_state, "samples_to_tensorboard") + + # table + frame = ctk.CTkFrame(master=master, corner_radius=0) + frame.grid(row=1, column=0, sticky="nsew") + + return SamplingTab(frame, self.train_config, self.ui_state) + + def create_backup_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + + # backup after + components.label(frame, 0, 0, "Backup After", + tooltip="The interval used when automatically creating model backups during training") + components.time_entry(frame, 0, 1, self.ui_state, "backup_after", "backup_after_unit") + + # backup now + components.button(frame, 0, 3, "backup now", self.backup_now) + + # rolling backup + components.label(frame, 1, 0, "Rolling Backup", + tooltip="If rolling backups are enabled, older backups are deleted automatically") + components.switch(frame, 1, 1, self.ui_state, "rolling_backup") + + # rolling backup count + components.label(frame, 1, 3, "Rolling Backup Count", + tooltip="Defines the number of backups to keep if rolling backups are enabled") + components.entry(frame, 1, 4, self.ui_state, "rolling_backup_count") + + # backup before save + components.label(frame, 2, 0, "Backup Before Save", + tooltip="Create a full backup before saving the final model") + components.switch(frame, 2, 1, self.ui_state, "backup_before_save") + + # save after + components.label(frame, 3, 0, "Save Every", + tooltip="The interval used when automatically saving the model during training") + components.time_entry(frame, 3, 1, self.ui_state, "save_every", "save_every_unit") + + # save now + components.button(frame, 3, 3, "save now", self.save_now) + + # skip save + components.label(frame, 4, 0, "Skip First", + tooltip="Start saving automatically after this interval has elapsed") + components.entry(frame, 4, 1, self.ui_state, "save_skip_first", width=50, sticky="nw") + + # save filename prefix + components.label(frame, 5, 0, "Save Filename Prefix", + tooltip="The prefix for filenames used when saving the model during training") + components.entry(frame, 5, 1, self.ui_state, "save_filename_prefix") + + frame.pack(fill="both", expand=1) + return frame + + def embedding_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + + # 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.path_entry( + frame, 0, 1, self.ui_state, "embedding.model_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # token count + components.label(frame, 1, 0, "Token count", + tooltip="The token count used when creating a new embedding. Leave empty to auto detect from the initial embedding text.") + components.entry(frame, 1, 1, self.ui_state, "embedding.token_count") + + # initial embedding text + components.label(frame, 2, 0, "Initial embedding text", + tooltip="The initial embedding text used when creating a new embedding") + components.entry(frame, 2, 1, self.ui_state, "embedding.initial_embedding_text") + + # embedding weight dtype + components.label(frame, 3, 0, "Embedding Weight Data Type", + tooltip="The Embedding weight data type used for training. This can reduce memory consumption, but reduces precision") + components.options_kv(frame, 3, 1, [ + ("float32", DataType.FLOAT_32), + ("bfloat16", DataType.BFLOAT_16), + ], self.ui_state, "embedding_weight_dtype") + + # placeholder + components.label(frame, 4, 0, "Placeholder", + tooltip="The placeholder used when using the embedding in a prompt") + components.entry(frame, 4, 1, self.ui_state, "embedding.placeholder") + + # output embedding + components.label(frame, 5, 0, "Output embedding", + tooltip="Output embeddings are calculated at the output of the text encoder, not the input. This can improve results for larger text encoders and lower VRAM usage.") + components.switch(frame, 5, 1, self.ui_state, "embedding.is_output_embedding") + + frame.pack(fill="both", expand=1) + return frame + + def create_additional_embeddings_tab(self, master): + return AdditionalEmbeddingsTab(master, self.train_config, self.ui_state) + + def create_tools_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + + # dataset + components.label(frame, 0, 0, "Dataset Tools", + tooltip="Open the captioning tool") + components.button(frame, 0, 1, "Open", self.open_dataset_tool) + + # video tools + components.label(frame, 1, 0, "Video Tools", + tooltip="Open the video tools") + components.button(frame, 1, 1, "Open", self.open_video_tool) + + # convert model + components.label(frame, 2, 0, "Convert Model Tools", + tooltip="Open the model conversion tool") + components.button(frame, 2, 1, "Open", self.open_convert_model_tool) + + # sample + components.label(frame, 3, 0, "Sampling Tool", + tooltip="Open the model sampling tool") + components.button(frame, 3, 1, "Open", self.open_sampling_tool) + + components.label(frame, 4, 0, "Profiling Tool", + tooltip="Open the profiling tools.") + components.button(frame, 4, 1, "Open", self.open_profiling_tool) + + frame.pack(fill="both", expand=1) + return frame + + def change_model_type(self, model_type: ModelType): + if self.model_tab: + self.model_tab.refresh_ui() + + if self.training_tab: + self.training_tab.refresh_ui() + + if self.lora_tab: + self.lora_tab.refresh_ui() + + def change_training_method(self, training_method: TrainingMethod): + if not self.tabview: + return + + if self.model_tab: + self.model_tab.refresh_ui() + + if training_method != TrainingMethod.LORA and "LoRA" in self.tabview._tab_dict: + self.tabview.delete("LoRA") + self.lora_tab = None + if training_method != TrainingMethod.EMBEDDING and "embedding" in self.tabview._tab_dict: + self.tabview.delete("embedding") + + if training_method == TrainingMethod.LORA and "LoRA" not in self.tabview._tab_dict: + self.lora_tab = LoraTab(self.tabview.add("LoRA"), self.train_config, self.ui_state) + if training_method == TrainingMethod.EMBEDDING and "embedding" not in self.tabview._tab_dict: + self.embedding_tab(self.tabview.add("embedding")) + + def load_preset(self): + if not self.tabview: + return + + if self.additional_embeddings_tab: + self.additional_embeddings_tab.refresh_ui() + + def open_tensorboard(self): + webbrowser.open("http://localhost:" + str(self.train_config.tensorboard_port), new=0, autoraise=False) + + def _calculate_eta_string(self, train_progress: TrainProgress, max_step: int, max_epoch: int) -> str | None: + spent_total = time.monotonic() - self.start_time + steps_done = train_progress.epoch * max_step + train_progress.epoch_step + remaining_steps = (max_epoch - train_progress.epoch - 1) * max_step + (max_step - train_progress.epoch_step) + total_eta = spent_total / steps_done * remaining_steps + + if train_progress.global_step <= 30: + return "Estimating ..." + + td = datetime.timedelta(seconds=total_eta) + days = td.days + hours, remainder = divmod(td.seconds, 3600) + minutes, seconds = divmod(remainder, 60) + if days > 0: + return f"{days}d {hours}h" + elif hours > 0: + return f"{hours}h {minutes}m" + elif minutes > 0: + return f"{minutes}m {seconds}s" + else: + return f"{seconds}s" + + def set_eta_label(self, train_progress: TrainProgress, max_step: int, max_epoch: int): + eta_str = self._calculate_eta_string(train_progress, max_step, max_epoch) + if eta_str is not None: + self.eta_label.configure(text=f"ETA: {eta_str}") + else: + self.eta_label.configure(text="") + + def delete_eta_label(self): + self.eta_label.configure(text="") + + def on_update_train_progress(self, train_progress: TrainProgress, max_step: int, max_epoch: int): + self.set_step_progress(train_progress.epoch_step, max_step) + self.set_epoch_progress(train_progress.epoch, max_epoch) + self.set_eta_label(train_progress, max_step, max_epoch) + + def on_update_status(self, status: str): + self.status_label.configure(text=status) + + def open_dataset_tool(self): + window = CaptionUI(self, None, False) + self.wait_window(window) + + def open_video_tool(self): + window = VideoToolUI(self) + self.wait_window(window) + + def open_convert_model_tool(self): + window = ConvertModelUI(self) + self.wait_window(window) + + def open_sampling_tool(self): + if not self.training_callbacks and not self.training_commands: + window = SampleWindow( + self, + use_external_model=False, + train_config=self.train_config, + ) + self.wait_window(window) + torch_gc() + + def open_profiling_tool(self): + self.profiling_window.deiconify() + + def generate_debug_package(self): + zip_path = filedialog.askdirectory( + initialdir=".", + title="Select Directory to Save Debug Package" + ) + + if not zip_path: + return + + zip_path = Path(zip_path) / "OneTrainer_debug_report.zip" + + self.on_update_status("Generating debug package...") + + try: + config_json_string = json.dumps(self.train_config.to_pack_dict(secrets=False)) + scripts.generate_debug_report.create_debug_package(str(zip_path), config_json_string) + self.on_update_status(f"Debug package saved to {zip_path.name}") + except Exception as e: + traceback.print_exc() + self.on_update_status(f"Error generating debug package: {e}") + + + def open_manual_sample_window (self): + training_callbacks = self.training_callbacks + training_commands = self.training_commands + + if training_callbacks and training_commands: + window = SampleWindow( + self, + train_config=self.train_config, + use_external_model=True, + callbacks=training_callbacks, + commands=training_commands, + ) + self.wait_window(window) + training_callbacks.set_on_sample_custom() + + def __training_thread_function(self): + error_caught = False + + self.training_callbacks = TrainCallbacks( + on_update_train_progress=self.on_update_train_progress, + on_update_status=self.on_update_status, + ) + + trainer = create.create_trainer(self.train_config, self.training_callbacks, self.training_commands, reattach=self.cloud_tab.reattach) + try: + trainer.start() + if self.train_config.cloud.enabled: + self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) + + self.start_time = time.monotonic() + trainer.train() + except Exception: + if self.train_config.cloud.enabled: + self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) + error_caught = True + traceback.print_exc() + + trainer.end() + + # clear gpu memory + del trainer + + 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() + + # 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) + + def start_training(self): + if self.training_thread is None: + self.save_default() + + # --- pre-training validation gate --- + errors = flush_and_validate_all() + + if errors: + bullet_list = "\n".join(f"• {e}" for e in errors) + messagebox.showerror( + "Cannot Start Training", + f"Please fix the following errors before training:\n\n{bullet_list}", + ) + return + + 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() + + self.training_thread = threading.Thread(target=self.__training_thread_function) + self.training_thread.start() + else: + self._set_training_button_stopping() + self.on_update_status("Stopping ...") + self.training_commands.stop() + + 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() + + def export_training(self): + file_path = filedialog.asksaveasfilename(filetypes=[ + ("All Files", "*.*"), + ("json", "*.json"), + ], initialdir=".", initialfile="config.json") + + if file_path: + with open(file_path, "w") as f: + json.dump(self.train_config.to_pack_dict(secrets=False), f, indent=4) + + def sample_now(self): + train_commands = self.training_commands + if train_commands: + train_commands.sample_default() + + def backup_now(self): + train_commands = self.training_commands + if train_commands: + train_commands.backup() + + def save_now(self): + train_commands = self.training_commands + 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") + + 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", + ] + + if self.train_config.tensorboard_expose: + tensorboard_args.append("--bind_all") + + 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 _set_training_button_style(self, mode: str): + if not self.training_button: + return + style = self._TRAIN_BUTTON_STYLES.get(mode) + if not style: + return + self.training_button.configure(**style) + + def _set_training_button_idle(self): + self._set_training_button_style("idle") + + def _set_training_button_running(self): + self._set_training_button_style("running") + + def _set_training_button_stopping(self): + self._set_training_button_style("stopping") diff --git a/modules/ui/CtkTrainingTabView.py b/modules/ui/CtkTrainingTabView.py new file mode 100644 index 000000000..bcca11ae9 --- /dev/null +++ b/modules/ui/CtkTrainingTabView.py @@ -0,0 +1,856 @@ +from modules.ui.OffloadingWindow import OffloadingWindow +from modules.ui.OptimizerParamsWindow import OptimizerParamsWindow +from modules.ui.SchedulerParamsWindow import SchedulerParamsWindow +from modules.ui.TimestepDistributionWindow import TimestepDistributionWindow +from modules.util import create +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.DataType import DataType +from modules.util.enum.EMAMode import EMAMode +from modules.util.enum.GradientCheckpointingMethod import GradientCheckpointingMethod +from modules.util.enum.LearningRateScaler import LearningRateScaler +from modules.util.enum.LearningRateScheduler import LearningRateScheduler +from modules.util.enum.LossScaler import LossScaler +from modules.util.enum.LossWeight import LossWeight +from modules.util.enum.Optimizer import Optimizer +from modules.util.enum.TimestepDistribution import TimestepDistribution +from modules.util.optimizer_util import change_optimizer +from modules.util.ui import components +from modules.util.ui.UIState import UIState +from modules.util.ui.validation_helpers import check_range, validate_resolution + +import customtkinter as ctk + + +class TrainingTab: + + def __init__(self, master, train_config: TrainConfig, ui_state: UIState): + super().__init__() + + self.master = master + self.train_config = train_config + self.ui_state = ui_state + + master.grid_rowconfigure(0, weight=1) + master.grid_columnconfigure(0, weight=1) + + self.scroll_frame = None + + self.refresh_ui() + + def refresh_ui(self): + if self.scroll_frame: + self.scroll_frame.destroy() + + self.scroll_frame = ctk.CTkScrollableFrame(self.master, fg_color="transparent") + self.scroll_frame.grid(row=0, column=0, sticky="nsew") + + self.scroll_frame.grid_columnconfigure(0, weight=1) + self.scroll_frame.grid_columnconfigure(1, weight=1) + self.scroll_frame.grid_columnconfigure(2, weight=1) + + column_0 = ctk.CTkFrame(master=self.scroll_frame, corner_radius=0, fg_color="transparent") + column_0.grid(row=0, column=0, sticky="nsew") + column_0.grid_columnconfigure(0, weight=1) + + column_1 = ctk.CTkFrame(master=self.scroll_frame, corner_radius=0, fg_color="transparent") + column_1.grid(row=0, column=1, sticky="nsew") + column_1.grid_columnconfigure(0, weight=1) + + column_2 = ctk.CTkFrame(master=self.scroll_frame, corner_radius=0, fg_color="transparent") + column_2.grid(row=0, column=2, sticky="nsew") + column_2.grid_columnconfigure(0, weight=1) + + if self.train_config.model_type.is_stable_diffusion(): + self.__setup_stable_diffusion_ui(column_0, column_1, column_2) + if self.train_config.model_type.is_stable_diffusion_3(): + self.__setup_stable_diffusion_3_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_stable_diffusion_xl(): + self.__setup_stable_diffusion_xl_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_wuerstchen(): + self.__setup_wuerstchen_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_pixart(): + self.__setup_pixart_alpha_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_flux_1(): + self.__setup_flux_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_flux_2(): + self.__setup_flux_2_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_chroma(): + self.__setup_chroma_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_qwen(): + self.__setup_qwen_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_sana(): + self.__setup_sana_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_hunyuan_video(): + self.__setup_hunyuan_video_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_hi_dream(): + self.__setup_hi_dream_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_z_image(): + self.__setup_z_image_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_ernie(): + self.__setup_ernie_ui(column_0, column_1, column_2) + + + def __setup_stable_diffusion_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_frame(column_0, 1) + self.__create_embedding_frame(column_0, 2) + + self.__create_base2_frame(column_1, 0, supports_circular_padding=True) + self.__create_unet_frame(column_1, 1) + self.__create_noise_frame(column_1, 2, supports_generalized_offset_noise=True) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_stable_diffusion_3_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 3, i=3, supports_include=True) + self.__create_embedding_frame(column_0, 4) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1) + self.__create_noise_frame(column_1, 2) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_stable_diffusion_xl_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_n_frame(column_0, 1, i=1) + self.__create_text_encoder_n_frame(column_0, 2, i=2) + self.__create_embedding_frame(column_0, 3) + + self.__create_base2_frame(column_1, 0, supports_circular_padding=True) + self.__create_unet_frame(column_1, 1) + self.__create_noise_frame(column_1, 2, supports_generalized_offset_noise=True) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_wuerstchen_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_frame(column_0, 1) + self.__create_embedding_frame(column_0, 2) + + self.__create_base2_frame(column_1, 0, supports_circular_padding=True) + self.__create_prior_frame(column_1, 1) + self.__create_noise_frame(column_1, 2) + + self.__create_masked_frame(column_2, 0) + self.__create_loss_frame(column_2, 1) + self.__create_layer_frame(column_2, 2) + + def __setup_pixart_alpha_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_frame(column_0, 1) + self.__create_embedding_frame(column_0, 2) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1) + self.__create_noise_frame(column_1, 2) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2, supports_vb_loss=True) + self.__create_layer_frame(column_2, 3) + + def __setup_flux_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True, supports_sequence_length=True) + self.__create_embedding_frame(column_0, 4) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1, supports_guidance_scale=True) + self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_flux_2_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False, supports_sequence_length=True) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1, supports_guidance_scale=True, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_chroma_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_frame(column_0, 1) + self.__create_embedding_frame(column_0, 4) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_qwen_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_z_image_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_ernie_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_sana_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_frame(column_0, 1) + self.__create_embedding_frame(column_0, 2) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1) + self.__create_noise_frame(column_1, 2) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_hunyuan_video_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True) + self.__create_embedding_frame(column_0, 4) + + self.__create_base2_frame(column_1, 0, video_training_enabled=True) + self.__create_transformer_frame(column_1, 1, supports_guidance_scale=True) + self.__create_noise_frame(column_1, 2) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_hi_dream_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 3, i=3, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 4, i=4, supports_include=True, supports_layer_skip=False) + self.__create_embedding_frame(column_0, 5) + + self.__create_base2_frame(column_1, 0, video_training_enabled=True) + self.__create_transformer_frame(column_1, 1) + self.__create_noise_frame(column_1, 2) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __create_base_frame(self, master, row): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + + # optimizer + components.label(frame, 0, 0, "Optimizer", + tooltip="The type of optimizer") + components.options_adv(frame, 0, 1, [str(x) for x in list(Optimizer)], self.ui_state, "optimizer.optimizer", + command=self.__restore_optimizer_config, adv_command=self.__open_optimizer_params_window) + + # learning rate scheduler + # Wackiness will ensue when reloading configs if we don't check and clear this first. + if hasattr(self, "lr_scheduler_comp"): + delattr(self, "lr_scheduler_comp") + delattr(self, "lr_scheduler_adv_comp") + components.label(frame, 1, 0, "Learning Rate Scheduler", + tooltip="Learning rate scheduler that automatically changes the learning rate during training") + _, d = components.options_adv(frame, 1, 1, [str(x) for x in list(LearningRateScheduler)], self.ui_state, + "learning_rate_scheduler", command=self.__restore_scheduler_config, + adv_command=self.__open_scheduler_params_window) + self.lr_scheduler_comp = d['component'] + self.lr_scheduler_adv_comp = d['button_component'] + # Initial call requires the presence of self.lr_scheduler_adv_comp. + self.__restore_scheduler_config(self.ui_state.get_var("learning_rate_scheduler").get()) + + # learning rate + components.label(frame, 2, 0, "Learning Rate", + tooltip="The base learning rate") + components.entry(frame, 2, 1, self.ui_state, "learning_rate", required=True) + + # learning rate warmup steps + components.label(frame, 3, 0, "Learning Rate Warmup Steps", + tooltip="The number of steps it takes to gradually increase the learning rate from 0 to the specified learning rate. Values >1 are interpeted as a fixed number of steps, values <=1 are intepreted as a percentage of the total training steps (ex. 0.2 = 20% of the total step count)") + components.entry(frame, 3, 1, self.ui_state, "learning_rate_warmup_steps") + + # learning rate min factor + components.label(frame, 4, 0, "Learning Rate Min Factor", + tooltip="Unit = float. Method = percentage. For a factor of 0.1, the final LR will be 10% of the initial LR. If the initial LR is 1e-4, the final LR will be 1e-5.") + components.entry(frame, 4, 1, self.ui_state, "learning_rate_min_factor", + extra_validate=check_range(lower=0, upper=0.99, message="Learning rate min factor must be between 0 and 0.99")) + + # learning rate cycles + components.label(frame, 5, 0, "Learning Rate Cycles", + tooltip="The number of learning rate cycles. This is only applicable if the learning rate scheduler supports cycles") + components.entry(frame, 5, 1, self.ui_state, "learning_rate_cycles") + + # epochs + components.label(frame, 6, 0, "Epochs", + tooltip="The number of epochs for a full training run") + components.entry(frame, 6, 1, self.ui_state, "epochs", required=True) + + # batch size + components.label(frame, 7, 0, "Local Batch Size", + tooltip="The batch size of one training step. If you use multiple GPUs, this is the batch size of each GPU (local batch size).") + components.entry(frame, 7, 1, self.ui_state, "batch_size", required=True) + + # accumulation steps + components.label(frame, 8, 0, "Accumulation Steps", + tooltip="Number of accumulation steps. Increase this number to trade batch size for training speed") + components.entry(frame, 8, 1, self.ui_state, "gradient_accumulation_steps", required=True) + + # Learning Rate Scaler + components.label(frame, 9, 0, "Learning Rate Scaler", + tooltip="Selects the type of learning rate scaling to use during training. Functionally equated as: LR * SQRT(selection)") + components.options(frame, 9, 1, [str(x) for x in list(LearningRateScaler)], self.ui_state, + "learning_rate_scaler") + + # clip grad norm + components.label(frame, 10, 0, "Clip Grad Norm", + tooltip="Clips the gradient norm. Leave empty to disable gradient clipping.") + components.entry(frame, 10, 1, self.ui_state, "clip_grad_norm") + + def __create_base2_frame(self, master, row, video_training_enabled: bool=False, supports_circular_padding: bool=False): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + row = 0 + + # ema + components.label(frame, row, 0, "EMA", + tooltip="EMA averages the training progress over many steps, better preserving different concepts in big datasets") + components.options(frame, row, 1, [str(x) for x in list(EMAMode)], self.ui_state, "ema") + row += 1 + + # ema decay + components.label(frame, row, 0, "EMA Decay", + tooltip="Decay parameter of the EMA model. Higher numbers will average more steps. For datasets of hundreds or thousands of images, set this to 0.9999. For smaller datasets, set it to 0.999 or even 0.998") + components.entry(frame, row, 1, self.ui_state, "ema_decay", + extra_validate=check_range(lower=0.5, upper=1, + message="EMA decay must be between 0.5 and 1")) + row += 1 + + # ema update step interval + components.label(frame, row, 0, "EMA Update Step Interval", + tooltip="Number of steps between EMA update steps") + components.entry(frame, row, 1, self.ui_state, "ema_update_step_interval") + row += 1 + + # gradient checkpointing + components.label(frame, row, 0, "Gradient checkpointing", + tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") + components.options_adv(frame, row, 1, [str(x) for x in list(GradientCheckpointingMethod)], self.ui_state, + "gradient_checkpointing", adv_command=self.__open_offloading_window) + row += 1 + + # gradient checkpointing layer offloading + components.label(frame, row, 0, "Layer offload fraction", + tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") + components.entry(frame, row, 1, self.ui_state, "layer_offload_fraction") + row += 1 + + # train dtype + components.label(frame, row, 0, "Train Data Type", + tooltip="The mixed precision data type used for training. This can increase training speed, but reduces precision") + components.options_kv(frame, row, 1, [ + ("float32", DataType.FLOAT_32), + ("float16", DataType.FLOAT_16), + ("bfloat16", DataType.BFLOAT_16), + ("tfloat32", DataType.TFLOAT_32), + ], self.ui_state, "train_dtype") + row += 1 + + # fallback train dtype + components.label(frame, row, 0, "Fallback Train Data Type", + tooltip="The mixed precision data type used for training stages that don't support float16 data types. This can increase training speed, but reduces precision") + components.options_kv(frame, row, 1, [ + ("float32", DataType.FLOAT_32), + ("bfloat16", DataType.BFLOAT_16), + ], self.ui_state, "fallback_train_dtype") + row += 1 + + # autocast cache + components.label(frame, row, 0, "Autocast Cache", + tooltip="Enables the autocast cache. Disabling this reduces memory usage, but increases training time") + components.switch(frame, row, 1, self.ui_state, "enable_autocast_cache") + row += 1 + + # resolution + components.label(frame, row, 0, "Resolution", + tooltip="The resolution used for training. Optionally specify multiple resolutions separated by a comma, or a single exact resolution in the format x") + components.entry(frame, row, 1, self.ui_state, "resolution", required=True, + extra_validate=validate_resolution()) + row += 1 + + # frames + if video_training_enabled: + components.label(frame, row, 0, "Frames", + tooltip="The number of frames used for training.") + components.entry(frame, row, 1, self.ui_state, "frames", required=True) + row += 1 + + # force circular padding + if supports_circular_padding: + components.label(frame, row, 0, "Force Circular Padding", + tooltip="Enables circular padding for all conv layers to better train seamless images") + components.switch(frame, row, 1, self.ui_state, "force_circular_padding") + + def __create_text_encoder_frame(self, master, row, supports_clip_skip=True, supports_training=True, supports_sequence_length=False): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + + if supports_training: + components.label(frame, 0, 0, "Train Text Encoder", + tooltip="Enables training the text encoder model") + components.switch(frame, 0, 1, self.ui_state, "text_encoder.train") + + # dropout + components.label(frame, 1, 0, "Caption Dropout Probability", + tooltip="The Probability for dropping the text encoder conditioning") + components.entry(frame, 1, 1, self.ui_state, "text_encoder.dropout_probability") + + if supports_training: + # train text encoder epochs + components.label(frame, 2, 0, "Stop Training After", + tooltip="When to stop training the text encoder") + components.time_entry(frame, 2, 1, self.ui_state, "text_encoder.stop_training_after", + "text_encoder.stop_training_after_unit", supports_time_units=False) + + # text encoder learning rate + components.label(frame, 3, 0, "Text Encoder Learning Rate", + tooltip="The learning rate of the text encoder. Overrides the base learning rate") + components.entry(frame, 3, 1, self.ui_state, "text_encoder.learning_rate") + + if supports_clip_skip: + # text encoder layer skip (clip skip) + components.label(frame, 4, 0, "Clip Skip", + tooltip="The number of additional clip layers to skip. 0 = the model default") + components.entry(frame, 4, 1, self.ui_state, "text_encoder_layer_skip") + + if supports_sequence_length: + # text encoder sequence length + components.label(frame, row, 0, "Text Encoder Sequence Length", + tooltip="Number of tokens for captions") + components.entry(frame, row, 1, self.ui_state, "text_encoder_sequence_length") + row += 1 + + def __create_text_encoder_n_frame( + self, + master, + row: int, + i: int, + supports_include: bool = False, + supports_layer_skip: bool = True, + supports_sequence_length: bool = False, + ): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + row = 0 + + suffix = f"_{i}" if i > 1 else "" + + if supports_include: + # include text encoder + components.label(frame, row, 0, f"Include Text Encoder {i}", + tooltip=f"Includes text encoder {i} in the training run") + components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.include") + row += 1 + + # train text encoder + components.label(frame, row, 0, f"Train Text Encoder {i}", + tooltip=f"Enables training the text encoder {i} model") + components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.train") + row += 1 + + # train text encoder embedding + components.label(frame, row, 0, f"Train Text Encoder {i} Embedding", + tooltip=f"Enables training embeddings for the text encoder {i} model") + components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.train_embedding") + row += 1 + + # dropout + components.label(frame, row, 0, "Dropout Probability", + tooltip=f"The Probability for dropping the text encoder {i} conditioning") + components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}.dropout_probability") + row += 1 + + # train text encoder epochs + components.label(frame, row, 0, "Stop Training After", + tooltip=f"When to stop training the text encoder {i}") + components.time_entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}.stop_training_after", + f"text_encoder{suffix}.stop_training_after_unit", supports_time_units=False) + row += 1 + + # text encoder learning rate + components.label(frame, row, 0, f"Text Encoder {i} Learning Rate", + tooltip=f"The learning rate of the text encoder {i}. Overrides the base learning rate") + components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}.learning_rate") + row += 1 + + if supports_layer_skip: + # text encoder layer skip (clip skip) + components.label(frame, row, 0, f"Text Encoder {i} Clip Skip", + tooltip="The number of additional clip layers to skip. 0 = the model default") + components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}_layer_skip") + row += 1 + + if supports_sequence_length: + # text encoder sequence length + components.label(frame, row, 0, f"Text Encoder {i} Sequence Length", + tooltip="Overrides the number of tokens used for captions. If empty, the model default is used, which is 512 on Flux. Comfy samples with 256 tokens though. 77 is the default only for backwards compatibility.") + components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}_sequence_length") + row += 1 + + def __create_embedding_frame(self, master, row): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + + # embedding learning rate + components.label(frame, 0, 0, "Embeddings Learning Rate", + tooltip="The learning rate of embeddings. Overrides the base learning rate") + components.entry(frame, 0, 1, self.ui_state, "embedding_learning_rate") + + # preserve embedding norm + components.label(frame, 1, 0, "Preserve Embedding Norm", + tooltip="Rescales each trained embedding to the median embedding norm") + components.switch(frame, 1, 1, self.ui_state, "preserve_embedding_norm") + + def __create_unet_frame(self, master, row): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + + # train unet + components.label(frame, 0, 0, "Train UNet", + tooltip="Enables training the UNet model") + components.switch(frame, 0, 1, self.ui_state, "unet.train") + + # train unet epochs + components.label(frame, 1, 0, "Stop Training After", + tooltip="When to stop training the UNet") + components.time_entry(frame, 1, 1, self.ui_state, "unet.stop_training_after", "unet.stop_training_after_unit", + supports_time_units=False) + + # unet learning rate + components.label(frame, 2, 0, "UNet Learning Rate", + tooltip="The learning rate of the UNet. Overrides the base learning rate") + components.entry(frame, 2, 1, self.ui_state, "unet.learning_rate") + + # rescale noise scheduler to zero terminal SNR + rescale_label = components.label(frame, 3, 0, "Rescale Noise Scheduler + V-pred", + tooltip="Rescales the noise scheduler to a zero terminal signal to noise ratio and switches the model to a v-prediction target") + rescale_label.configure(wraplength=130, justify="left") + components.switch(frame, 3, 1, self.ui_state, "rescale_noise_scheduler_to_zero_terminal_snr") + + def __create_prior_frame(self, master, row): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + + # train prior + components.label(frame, 0, 0, "Train Prior", + tooltip="Enables training the Prior model") + components.switch(frame, 0, 1, self.ui_state, "prior.train") + + # train prior epochs + components.label(frame, 1, 0, "Stop Training After", + tooltip="When to stop training the Prior") + components.time_entry(frame, 1, 1, self.ui_state, "prior.stop_training_after", "prior.stop_training_after_unit", + supports_time_units=False) + + # prior learning rate + components.label(frame, 2, 0, "Prior Learning Rate", + tooltip="The learning rate of the Prior. Overrides the base learning rate") + components.entry(frame, 2, 1, self.ui_state, "prior.learning_rate") + + def __create_transformer_frame(self, master, row, supports_guidance_scale: bool = False, supports_force_attention_mask: bool = True): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + + # train transformer + components.label(frame, 0, 0, "Train Transformer", + tooltip="Enables training the Transformer model") + components.switch(frame, 0, 1, self.ui_state, "transformer.train") + + # train transformer epochs + components.label(frame, 1, 0, "Stop Training After", + tooltip="When to stop training the Transformer") + components.time_entry(frame, 1, 1, self.ui_state, "transformer.stop_training_after", "transformer.stop_training_after_unit", + supports_time_units=False) + + # transformer learning rate + components.label(frame, 2, 0, "Transformer Learning Rate", + tooltip="The learning rate of the Transformer. Overrides the base learning rate") + components.entry(frame, 2, 1, self.ui_state, "transformer.learning_rate") + + if supports_force_attention_mask: + # transformer learning rate + components.label(frame, 3, 0, "Force Attention Mask", + tooltip="Force enables passing of a text embedding attention mask to the transformer. This can improve training on shorter captions.") + components.switch(frame, 3, 1, self.ui_state, "transformer.attention_mask") + + if supports_guidance_scale: + # guidance scale + components.label(frame, 4, 0, "Guidance Scale", + tooltip="The guidance scale of guidance distilled models passed to the transformer during training.") + components.entry(frame, 4, 1, self.ui_state, "transformer.guidance_scale") + + def __create_noise_frame(self, master, row, supports_generalized_offset_noise: bool = False, supports_dynamic_timestep_shifting: bool = False): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + + # offset noise weight + components.label(frame, 0, 0, "Offset Noise Weight", + tooltip="The weight of offset noise added to each training step") + components.entry(frame, 0, 1, self.ui_state, "offset_noise_weight") + + if supports_generalized_offset_noise: + # generalized offset noise weight + generalised_offset_label = components.label(frame, 1, 0, "Generalized Offset Noise", + tooltip="Per-timestep 'brightness knob' instead of a fixed offset - steadier training, better starts, and improved very dark/bright images. Compatible with V-pred and Eps-pred. Start with 0.02 and adjust as needed.") + generalised_offset_label.configure(wraplength=130, justify="left") + components.switch(frame, 1, 1, self.ui_state, "generalized_offset_noise") + + # perturbation noise weight + components.label(frame, 2, 0, "Perturbation Noise Weight", + tooltip="The weight of perturbation noise added to each training step") + components.entry(frame, 2, 1, self.ui_state, "perturbation_noise_weight") + + # timestep distribution + components.label(frame, 3, 0, "Timestep Distribution", + tooltip="Selects the function to sample timesteps during training", + wide_tooltip=True) + components.options_adv(frame, 3, 1, [str(x) for x in list(TimestepDistribution)], self.ui_state, "timestep_distribution", + adv_command=self.__open_timestep_distribution_window) + + # min noising strength + components.label(frame, 4, 0, "Min Noising Strength", + tooltip="Specifies the minimum noising strength used during training. This can help to improve composition, but prevents finer details from being trained") + components.entry(frame, 4, 1, self.ui_state, "min_noising_strength", required=True) + + # max noising strength + components.label(frame, 5, 0, "Max Noising Strength", + tooltip="Specifies the maximum noising strength used during training. This can be useful to reduce overfitting, but also reduces the impact of training samples on the overall image composition") + components.entry(frame, 5, 1, self.ui_state, "max_noising_strength", required=True) + + # noising weight + components.label(frame, 6, 0, "Noising Weight", + tooltip="Controls the weight parameter of the timestep distribution function. Use the preview to see more details.") + components.entry(frame, 6, 1, self.ui_state, "noising_weight", required=True) + + # noising bias + components.label(frame, 7, 0, "Noising Bias", + tooltip="Controls the bias parameter of the timestep distribution function. Use the preview to see more details.") + components.entry(frame, 7, 1, self.ui_state, "noising_bias", required=True) + + # timestep shift + components.label(frame, 8, 0, "Timestep Shift", + tooltip="Shift the timestep distribution. Use the preview to see more details.") + components.entry(frame, 8, 1, self.ui_state, "timestep_shift", required=True) + + if supports_dynamic_timestep_shifting: + # dynamic timestep shifting + components.label(frame, 9, 0, "Dynamic Timestep Shifting", + tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) + components.switch(frame, 9, 1, self.ui_state, "dynamic_timestep_shifting") + + + + def __create_masked_frame(self, master, row): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + + # Masked Training + components.label(frame, 0, 0, "Masked Training", + tooltip="Masks the training samples to let the model focus on certain parts of the image. When enabled, one mask image is loaded for each training sample.") + components.switch(frame, 0, 1, self.ui_state, "masked_training") + + # unmasked probability + components.label(frame, 1, 0, "Unmasked Probability", + tooltip="When masked training is enabled, specifies the number of training steps done on unmasked samples") + components.entry(frame, 1, 1, self.ui_state, "unmasked_probability", + extra_validate=check_range(lower=0, upper=1, message="Unmasked probability must be between 0 and 1")) + + # unmasked weight + components.label(frame, 2, 0, "Unmasked Weight", + tooltip="When masked training is enabled, specifies the loss weight of areas outside the masked region") + components.entry(frame, 2, 1, self.ui_state, "unmasked_weight", + extra_validate=check_range(lower=0, upper=1, message="Unmasked weight must be between 0 and 1")) + + # normalize masked area loss + components.label(frame, 3, 0, "Normalize Masked Area Loss", + tooltip="When masked training is enabled, normalizes the loss for each sample based on the sizes of the masked region") + components.switch(frame, 3, 1, self.ui_state, "normalize_masked_area_loss") + + # masked prior preservation + components.label(frame, 4, 0, "Masked Prior Preservation Weight", + tooltip="Preserves regions outside the mask using the original untrained model output as a target. Only available for LoRA training. If enabled, use a low unmasked weight.") + components.entry(frame, 4, 1, self.ui_state, "masked_prior_preservation_weight", + extra_validate=check_range(lower=0, upper=1, message="Masked prior preservation weight must be between 0 and 1")) + + # use custom conditioning image + components.label(frame, 5, 0, "Custom Conditioning Image", + tooltip="When custom conditioning image is enabled, will use png postfix with -condlabel instead of automatically generated.It's suitable for special scenarios, such as object removal, allowing the model to learn a certain behavior concept") + components.switch(frame, 5, 1, self.ui_state, "custom_conditioning_image") + + def __create_loss_frame(self, master, row, supports_vb_loss: bool = False): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + + # MSE Strength + components.label(frame, 0, 0, "MSE Strength", + tooltip="Mean Squared Error strength for custom loss settings. Strengths should generally sum to 1.") + components.entry(frame, 0, 1, self.ui_state, "mse_strength", required=True) + + # MAE Strength + components.label(frame, 1, 0, "MAE Strength", + tooltip="Mean Absolute Error strength for custom loss settings. Strengths should generally sum to 1.") + components.entry(frame, 1, 1, self.ui_state, "mae_strength", required=True) + + # log-cosh Strength + components.label(frame, 2, 0, "log-cosh Strength", + tooltip="Log - Hyperbolic cosine Error strength for custom loss settings. Strengths should generally sum to 1.") + components.entry(frame, 2, 1, self.ui_state, "log_cosh_strength", required=True) + + # Huber Strength + components.label(frame, 3, 0, "Huber Strength", + tooltip="Huber loss strength for custom loss settings. Less sensitive to outliers than MSE. Strengths should generally sum to 1.") + components.entry(frame, 3, 1, self.ui_state, "huber_strength", required=True) + + # Huber Delta + components.label(frame, 4, 0, "Huber Delta", + tooltip="Delta parameter for huber loss") + components.entry(frame, 4, 1, self.ui_state, "huber_delta", required=True) + + if supports_vb_loss: + # VB Strength + components.label(frame, 5, 0, "VB Strength", + tooltip="Variational lower-bound strength for custom loss settings. Should be set to 1 for variational diffusion models") + components.entry(frame, 5, 1, self.ui_state, "vb_loss_strength", required=True) + + # Loss Weight function + components.label(frame, 6, 0, "Loss Weight Function", + tooltip="Choice of loss weight function. Can help the model learn details more accurately.") + components.options(frame, 6, 1, [str(x) for x in list(LossWeight) + if x.supports_flow_matching() == self.train_config.model_type.is_flow_matching() + or x == LossWeight.CONSTANT + ], + self.ui_state, "loss_weight_fn") + + row = 7 + + # Loss weight strength + if not self.train_config.model_type.is_flow_matching(): + components.label(frame, row, 0, "Gamma", + tooltip="Inverse strength of loss weighting. Range: 1-20, only applies to Min SNR and P2.") + components.entry(frame, row, 1, self.ui_state, "loss_weight_strength", + extra_validate=check_range(lower=1, upper=20, message="Gamma must be between 1 and 20")) + row += 1 + + # Loss Scaler + components.label(frame, row, 0, "Loss Scaler", + tooltip="Selects the type of loss scaling to use during training. Functionally equated as: Loss * selection") + components.options(frame, row, 1, [str(x) for x in list(LossScaler)], self.ui_state, "loss_scaler") + row += 1 + + def __create_layer_frame(self, master, row): + cls = create.get_model_setup_class(self.train_config.model_type, self.train_config.training_method) + presets = cls.LAYER_PRESETS if cls is not None else {"full": []} + components.layer_filter_entry(master, row, 0, self.ui_state, + preset_var_name="layer_filter_preset", presets=presets, + preset_label="Layer Filter", + preset_tooltip="Select a preset defining which layers to train, or select 'Custom' to define your own.\nA blank 'custom' field or 'Full' will train all layers.", + entry_var_name="layer_filter", + entry_tooltip="Comma-separated list of diffusion layers to train. Regular expressions (if toggled) are supported. Any model layer with a matching name will be trained", + regex_var_name="layer_filter_regex", + regex_tooltip="If enabled, layer filter patterns are interpreted as regular expressions. Otherwise, simple substring matching is used.", + ) + + + def __on_layer_filter_preset_change(self): + if not self.layer_selector: + return + selected = self.ui_state.get_var("layer_filter_preset").get() + self.__preset_set_layer_choice(selected) + + def __hide_layer_entry(self): + if self.layer_entry and self.layer_entry.winfo_manager(): + self.layer_entry.grid_remove() + + def __show_layer_entry(self): + if self.layer_entry and not self.layer_entry.winfo_manager(): + self.layer_entry.grid() + + def __open_optimizer_params_window(self): + window = OptimizerParamsWindow(self.master, self.train_config, self.ui_state) + self.master.wait_window(window) + + def __open_scheduler_params_window(self): + window = SchedulerParamsWindow(self.master, self.train_config, self.ui_state) + self.master.wait_window(window) + + def __open_timestep_distribution_window(self): + window = TimestepDistributionWindow(self.master, self.train_config, self.ui_state) + self.master.wait_window(window) + + def __open_offloading_window(self): + window = OffloadingWindow(self.master, self.train_config, self.ui_state) + self.master.wait_window(window) + + def __restore_optimizer_config(self, *args): + optimizer_config = change_optimizer(self.train_config) + self.ui_state.get_var("optimizer").update(optimizer_config) + + def __restore_scheduler_config(self, variable): + if not hasattr(self, 'lr_scheduler_adv_comp'): + return + + if variable == "CUSTOM": + self.lr_scheduler_adv_comp.configure(state="normal") + else: + self.lr_scheduler_adv_comp.configure(state="disabled") diff --git a/modules/ui/CtkVideoToolUIView.py b/modules/ui/CtkVideoToolUIView.py new file mode 100644 index 000000000..c3291e6ea --- /dev/null +++ b/modules/ui/CtkVideoToolUIView.py @@ -0,0 +1,877 @@ +import concurrent.futures +import math +import os +import pathlib +import random +import shlex +import subprocess +import threading +import webbrowser +from fractions import Fraction +from tkinter import filedialog + +from modules.util.image_util import load_image +from modules.util.path_util import SUPPORTED_VIDEO_EXTENSIONS +from modules.util.ui import components + +import av +import customtkinter as ctk +import cv2 +import scenedetect +from PIL import Image + + +class VideoToolUI(ctk.CTkToplevel): + def __init__( + self, + parent, + *args, **kwargs, + ): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + + self.title("Video Tools") + self.geometry("600x720") + self.resizable(True, True) + self.wait_visibility() + self.focus_set() + + self.grid_rowconfigure(0, weight=1) + self.grid_rowconfigure(1, weight=0) + self.grid_columnconfigure(0, weight=1) + + tabview = ctk.CTkTabview(self) + tabview.grid(row=0, column=0, sticky="nsew") + + self.clip_extract_tab = self.__clip_extract_tab(tabview.add("extract clips")) + self.image_extract_tab = self.__image_extract_tab(tabview.add("extract images")) + self.video_download_tab = self.__video_download_tab(tabview.add("download")) + self.status_bar(self) + + def status_bar(self, master): + frame = ctk.CTkFrame(master, fg_color="transparent") + frame.grid(row=1, column=0) + frame.grid_columnconfigure(0, weight=0, minsize=160) + frame.grid_columnconfigure(1, weight=0, minsize=300) + frame.grid_columnconfigure(2, weight=1) + + #create preview image + preview_path = "resources/icons/icon.png" + preview = load_image(preview_path, 'RGB') + preview.thumbnail((150, 150)) + self.preview_image= ctk.CTkImage(light_image=preview, size=preview.size) + self.preview_image_label = ctk.CTkLabel( + master=frame, text="Preview image", image=self.preview_image, height=150, width=150, + compound="top") + self.preview_image_label.grid(row=0, column=0, sticky="nw", padx=5, pady=5) + + #displays progress and messages that also go to terminal + self.status_label = ctk.CTkTextbox(master=frame, width=400, height=160, wrap="word", border_width=2) + self.status_label.insert(index="1.0", text="Current status") + self.status_label.configure(state="disabled") + self.status_label.grid(row=0, column=1, sticky="ne", padx=5, pady=5) + + def __clip_extract_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0, minsize=120) + frame.grid_columnconfigure(1, weight=0, minsize=200) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # single video + components.label(frame, 0, 0, "Single Video", + tooltip="Link to single video file to process.") + self.clip_single_entry = ctk.CTkEntry(frame, width=190) + self.clip_single_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) + self.clip_single_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_file(self.clip_single_entry, + [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))] + )) + self.clip_single_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) + components.button(frame, 0, 2, "Extract Single", + command=lambda: self.__extract_clips_button(False)) + + # time range + components.label(frame, 1, 0, " Time Range", + tooltip="Time range to limit selection for single video, \ + format as hour:minute:second, minute:second, or seconds.") + self.clip_time_start_entry = ctk.CTkEntry(frame, width=100) + self.clip_time_start_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.clip_time_start_entry.insert(0, "00:00:00") + self.clip_time_end_entry = ctk.CTkEntry(frame, width=100) + self.clip_time_end_entry.grid(row=1, column=1, sticky="e", padx=5, pady=5) + self.clip_time_end_entry.insert(0, "99:99:99") + + # directory of videos + components.label(frame, 2, 0, "Directory", + tooltip="Path to directory with multiple videos to process, including in subdirectories.") + self.clip_list_entry = ctk.CTkEntry(frame, width=190) + self.clip_list_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + self.clip_list_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_dir(self.clip_list_entry)) + self.clip_list_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) + components.button(frame, 2, 2, "Extract Directory", + command=lambda: self.__extract_clips_button(True)) + + # output directory + components.label(frame, 3, 0, "Output", + tooltip="Path to folder where extracted clips will be saved.") + self.clip_output_entry = ctk.CTkEntry(frame, width=190) + self.clip_output_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) + self.clip_output_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_dir(self.clip_output_entry)) + self.clip_output_button.grid(row=3, column=1, sticky="e", padx=5, pady=5) + + # output to subdirectories + self.output_subdir_clip = ctk.BooleanVar(self, False) + components.label(frame, 4, 0, "Output to\nSubdirectories", + tooltip="If enabled, files are saved to subfolders based on filename and input directory. \ + Otherwise will all be saved to the top level of the output directory.") + self.output_subdir_clip_entry = ctk.CTkSwitch(frame, variable=self.output_subdir_clip, text="") + self.output_subdir_clip_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + + # split at cuts + self.split_at_cuts = ctk.BooleanVar(self, False) + components.label(frame, 5, 0, "Split at Cuts", + tooltip="If enabled, detect cuts in the input video and split at those points. \ + Otherwise will split at any point, and clips may contain cuts.") + self.split_cuts_entry = ctk.CTkSwitch(frame, variable=self.split_at_cuts, text="") + self.split_cuts_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) + + # maximum length + components.label(frame, 6, 0, "Max Length (s)", + tooltip="Maximum length in seconds for saved clips, larger clips will be broken into multiple small clips.") + self.clip_length_entry = ctk.CTkEntry(frame, width=220) + self.clip_length_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) + self.clip_length_entry.insert(0, "3") + + # Set FPS + components.label(frame, 7, 0, "Set FPS", + tooltip="FPS to convert output videos to, set to 0 to keep original rate.") + self.clip_fps_entry = ctk.CTkEntry(frame, width=220) + self.clip_fps_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) + self.clip_fps_entry.insert(0, "24.0") + + # Remove borders + self.clip_bordercrop = ctk.BooleanVar(self, False) + components.label(frame, 8, 0, "Remove Borders", + tooltip="Remove black borders from output clip") + self.clip_bordercrop_entry = ctk.CTkSwitch(frame, variable=self.clip_bordercrop, text="") + self.clip_bordercrop_entry.grid(row=8, column=1, sticky="w", padx=5, pady=5) + + # Crop Variation + components.label(frame, 9, 0, "Crop Variation", + tooltip="Output clips will be randomly cropped to +- the base aspect ratio, \ + somewhat biased towards making square videos. Set to 0 to use only base aspect.") + self.clip_crop_entry = ctk.CTkEntry(frame, width=220) + self.clip_crop_entry.grid(row=9, column=1, sticky="w", padx=5, pady=5) + self.clip_crop_entry.insert(0, "0.2") + + # object filter - currently unused, may implement in future + # components.label(frame, 9, 0, "Object Filter", + # tooltip="Detect general features using Haar-Cascade classifier, and choose how to deal with clips where it is detected") + # components.options(frame, 9, 1, ["NONE", "FACE", "EYE", "BODY"], self.video_ui_state, "filter_object") + # components.options(frame, 9, 2, ["INCLUDE", "EXCLUDE", "SUBFOLDER"], self.video_ui_state, "filter_behavior") + + frame.pack(fill="both", expand=1) + return frame + + def __image_extract_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0, minsize=120) + frame.grid_columnconfigure(1, weight=0, minsize=200) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # single video + components.label(frame, 0, 0, "Single Video", + tooltip="Link to single video file to process.") + self.image_single_entry = ctk.CTkEntry(frame, width=190) + self.image_single_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) + self.image_single_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_file(self.image_single_entry, + [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))] + )) + self.image_single_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) + components.button(frame, 0, 2, "Extract Single", + command=lambda: self.__extract_images_button(False)) + + # time range + components.label(frame, 1, 0, " Time Range", + tooltip="Time range to limit selection for single video, \ + format as hour:minute:second, minute:second, or seconds.") + self.image_time_start_entry = ctk.CTkEntry(frame, width=100) + self.image_time_start_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.image_time_start_entry.insert(0, "00:00:00") + self.image_time_end_entry = ctk.CTkEntry(frame, width=100) + self.image_time_end_entry.grid(row=1, column=1, sticky="e", padx=5, pady=5) + self.image_time_end_entry.insert(0, "99:99:99") + + # directory of videos + components.label(frame, 2, 0, "Directory", + tooltip="Path to directory with multiple videos to process, including in subdirectories.") + self.image_list_entry = ctk.CTkEntry(frame, width=190) + self.image_list_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + self.image_list_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_dir(self.image_list_entry)) + self.image_list_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) + components.button(frame, 2, 2, "Extract Directory", + command=lambda: self.__extract_images_button(True)) + + # output directory + components.label(frame, 3, 0, "Output", + tooltip="Path to folder where extracted images will be saved.") + self.image_output_entry = ctk.CTkEntry(frame, width=190) + self.image_output_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) + self.image_output_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_dir(self.image_output_entry)) + self.image_output_button.grid(row=3, column=1, sticky="e", padx=5, pady=5) + + # output to subdirectories + self.output_subdir_img = ctk.BooleanVar(self, False) + components.label(frame, 4, 0, "Output to\nSubdirectories", + tooltip="If enabled, files are saved to subfolders based on filename and input directory. \ + Otherwise will all be saved to the top level of the output directory.") + self.output_subdir_img_entry = ctk.CTkSwitch(frame, variable=self.output_subdir_img, text="") + self.output_subdir_img_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + + # image capture rate + components.label(frame, 5, 0, "Images/sec", + tooltip="Number of images to capture per second of video. \ + Images will be taken at semi-random frames around the specified frequency.") + self.capture_rate_entry = ctk.CTkEntry(frame, width=220) + self.capture_rate_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) + self.capture_rate_entry.insert(0, "0.5") + + # blur removal + components.label(frame, 6, 0, "Blur Removal", + tooltip="Threshold for removal of blurry images, relative to all others. \ + For example at 0.2, the blurriest 20%% of the final selected frames will not be saved.") + self.blur_threshold_entry = ctk.CTkEntry(frame, width=220) + self.blur_threshold_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) + self.blur_threshold_entry.insert(0, "0.2") + + # Remove borders + self.image_bordercrop = ctk.BooleanVar(self, False) + components.label(frame, 7, 0, "Remove Borders", + tooltip="Remove black borders from output image") + self.image_bordercrop_entry = ctk.CTkSwitch(frame, variable=self.image_bordercrop, text="") + self.image_bordercrop_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) + + # Crop Variation + components.label(frame, 8, 0, "Crop Variation", + tooltip="Output images will be randomly cropped to +- the base aspect ratio, \ + somewhat biased towards making square images. Set to 0 to use only base sapect.") + self.image_crop_entry = ctk.CTkEntry(frame, width=220) + self.image_crop_entry.grid(row=8, column=1, sticky="w", padx=5, pady=5) + self.image_crop_entry.insert(0, "0.2") + + # # object filter - currently unused, may implement in future + # components.label(frame, 5, 0, "Object Filter", + # tooltip="Detect general features using Haar-Cascade classifier, and choose how to deal with clips where it is detected") + # components.options(frame, 5, 1, ["NONE", "FACE", "EYE", "BODY"], self.video_ui_state, "filter_object") + # components.options(frame, 5, 2, ["INCLUDE", "EXCLUDE", "SUBFOLDER"], self.video_ui_state, "filter_behavior") + + frame.pack(fill="both", expand=1) + return frame + + def __video_download_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0, minsize=120) + frame.grid_columnconfigure(1, weight=0, minsize=200) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # link + components.label(frame, 0, 0, "Single Link", + tooltip="Link to video/playlist to download. Uses yt-dlp, supports youtube, twitch, instagram, and many other sites.") + self.download_link_entry = ctk.CTkEntry(frame, width=220) + self.download_link_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) + components.button(frame, 0, 2, "Download Link", command=lambda: self.__download_button(False)) + + # link list + components.label(frame, 1, 0, "Link List", + tooltip="Path to txt file with list of links separated by newlines.") + self.download_list_entry = ctk.CTkEntry(frame, width=190) + self.download_list_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.download_list_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_file(self.download_list_entry, [("Text file", ".txt")])) + self.download_list_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) + components.button(frame, 1, 2, "Download List", command=lambda: self.__download_button(True)) + + # output directory + components.label(frame, 2, 0, "Output", + tooltip="Path to folder where downloaded videos will be saved.") + self.download_output_entry = ctk.CTkEntry(frame, width=190) + self.download_output_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + self.download_output_button = ctk.CTkButton(frame, width=30, text="...", command=lambda: self.__browse_for_dir(self.download_output_entry)) + self.download_output_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) + + # additional args + 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.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 --format mp4") + components.button(frame, 3, 2, "yt-dlp info", + command=lambda: webbrowser.open("https://github.com/yt-dlp/yt-dlp?tab=readme-ov-file#usage-and-options", new=0, autoraise=False)) + + frame.pack(fill="both", expand=1) + return frame + + def __browse_for_dir(self, entry_box): + # get the path from the user + path = filedialog.askdirectory() + # set the path to the entry box + # delete entry box text + entry_box.focus_set() + entry_box.delete(0, ctk.END) + entry_box.insert(0, path) + self.focus_set() + + def __browse_for_file(self, entry_box, filetypes): + # get the path from the user + path = filedialog.askopenfilename(filetypes=filetypes) + # set the path to the entry box + # delete entry box text + entry_box.focus_set() + entry_box.delete(0, ctk.END) + entry_box.insert(0, path) + self.focus_set() + + def __get_vid_paths(self, batch_mode: bool, input_path_single: str, input_path_dir: str): + input_videos = [] + if not batch_mode: + path = pathlib.Path(input_path_single) + if path.is_file(): + vid = cv2.VideoCapture(str(path)) + ok = False + try: + if vid.isOpened(): + ok, _ = vid.read() + finally: + vid.release() + if ok: + return [path] + else: + self.__update_status("Invalid video file!") + return [] + else: + self.__update_status("No file specified, or invalid file path!") + return [] + else: + input_videos = [] + if not pathlib.Path(input_path_dir).is_dir() or input_path_dir == "": + self.__update_status("Invalid input directory!") + return [] + # Only traverse supported extensions to avoid opening every file. + lower_exts = {e.lower() for e in SUPPORTED_VIDEO_EXTENSIONS} + for path in pathlib.Path(input_path_dir).rglob("*"): + if path.is_file() and path.suffix.lower() in lower_exts: + vid = cv2.VideoCapture(str(path)) + ok = False + try: + if vid.isOpened(): + ok, _ = vid.read() + finally: + vid.release() + if ok: + input_videos.append(path) + self.__update_status(f'Found {len(input_videos)} videos to process') + return input_videos + + def __run_in_thread(self, target, *args): + """Clear status box and run target function in a daemon thread.""" + self.status_label.configure(state="normal") + self.status_label.delete(index1="1.0", index2="end") + self.status_label.configure(state="disabled") + t = threading.Thread(target=target, args=args) + t.daemon = True + t.start() + + @staticmethod + def __parse_timestamp_to_frames(timestamp: str, fps: float) -> int: + return int(sum(int(x) * 60 ** i for i, x in enumerate(reversed(timestamp.split(':')))) * fps) + + def __get_safe_fps(self, video: cv2.VideoCapture, video_path: str) -> float: + fps = video.get(cv2.CAP_PROP_FPS) or 0.0 + if fps <= 0: + self.__update_status(f'Warning: Could not read FPS for "{os.path.basename(video_path)}". Falling back to 30 FPS.') + return 30.0 + return fps + + @staticmethod + def __get_output_dir(use_subdir: bool, batch_mode: bool, output_entry: str, + video_path, input_dir: str) -> str: + if use_subdir and batch_mode: + return os.path.join(output_entry, + os.path.splitext(os.path.relpath(video_path, input_dir))[0]) + elif use_subdir: + return os.path.join(output_entry, + os.path.splitext(os.path.basename(video_path))[0]) + return output_entry + + def __get_random_aspect(self, height: int, width: int, variation: float) -> tuple[int, int, int, int]: + # Return original dimensions and no offset if variation is zero + if variation == 0: + return 0, height, 0, width + + old_aspect = height/width + variation_scaled = old_aspect*variation + if old_aspect > 1.2: #tall image + new_aspect = min(4.0, max(1.0, random.triangular(old_aspect-(variation_scaled*1.5), old_aspect+(variation_scaled/2), old_aspect))) + elif old_aspect < 0.85: #wide image + new_aspect = max(0.25, min(1.0, random.triangular(old_aspect-(variation_scaled/2), old_aspect+(variation_scaled*1.5), old_aspect))) + else: #square image + new_aspect = random.triangular(old_aspect-variation_scaled, old_aspect+variation_scaled) + + new_aspect = round(new_aspect, 2) + #keep the height the same if reducing width, and vice versa + if new_aspect > old_aspect: + new_height = int(height) + new_width = int(width*(old_aspect/new_aspect)) + elif new_aspect < old_aspect: + new_height = int(height*(new_aspect/old_aspect)) + new_width = int(width) + else: + new_height = int(height) + new_width = int(width) + + #random offset in dimension that was cropped + position_x = random.randint(0, width-new_width) + position_y = random.randint(0, height-new_height) + return position_y, new_height, position_x, new_width + + def find_main_contour(self, frame): + #outline image to find main content and exclude black bars often present on letterboxed videos + frame_grayscale = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) + _, frame_thresh = cv2.threshold(frame_grayscale, 15, 255, cv2.THRESH_BINARY) + frame_contours, _ = cv2.findContours(frame_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + if frame_contours: + #select largest contour by area + frame_maincontour = max(frame_contours, key=lambda c: cv2.contourArea(c)) + x1, y1, w1, h1 = cv2.boundingRect(frame_maincontour) + else: #fallback if no contours detected + x1 = 0 + y1 = 0 + h1, w1, _ = frame.shape + + #if bounding box did not detect the correct area, likely due to all-black frame + if not frame_contours or h1 < 10 or w1 < 10: + x1 = 0 + y1 = 0 + h1, w1, _ = frame.shape + return x1, y1, w1, h1 + + def __extract_clips_button(self, batch_mode: bool): + self.__run_in_thread(self.__extract_clips_multi, batch_mode) + + def __extract_clips_multi(self, batch_mode: bool): + if not pathlib.Path(self.clip_output_entry.get()).is_dir() or self.clip_output_entry.get() == "": + self.__update_status("Invalid output directory!") + return + + # validate numeric inputs + try: + max_length = float(self.clip_length_entry.get()) + crop_variation = float(self.clip_crop_entry.get()) + target_fps = float(self.clip_fps_entry.get()) + input_single_entry = self.clip_single_entry.get() + input_multiple_entry = self.clip_list_entry.get() + output_entry = self.clip_output_entry.get() + except ValueError: + self.__update_status("Invalid numeric input for Max Length, Crop Variation, or FPS.") + return + if max_length <= 0.25: + self.__update_status("Max Length of clips must be > 0.25 seconds.") + return + if target_fps < 0: + self.__update_status("Target FPS must be a positive number (or 0 to skip fps re-encoding).") + return + if not (0.0 <= crop_variation < 1.0): + self.__update_status("Crop Variation must be between 0.0 and 1.0.") + return + + input_videos = self.__get_vid_paths(batch_mode, input_single_entry, input_multiple_entry) + if len(input_videos) == 0: # exit if no paths found + return + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + for video_path in input_videos: + output_directory = self.__get_output_dir( + self.output_subdir_clip_entry.get(), batch_mode, + output_entry, video_path, input_multiple_entry) + time_start = "00:00:00" if batch_mode else str(self.clip_time_start_entry.get()) + time_end = "99:99:99" if batch_mode else str(self.clip_time_end_entry.get()) + executor.submit(self.__extract_clips, + str(video_path), time_start, time_end, max_length, + self.split_at_cuts.get(), bool(self.clip_bordercrop_entry.get()), + crop_variation, target_fps, output_directory) + + if batch_mode: + self.__update_status(f'Clip extraction from all videos in "{input_multiple_entry}" complete') + else: + self.__update_status(f'Clip extraction from "{input_single_entry}" complete') + + def __extract_clips(self, video_path: str, timestamp_min: str, timestamp_max: str, max_length: float, + split_at_cuts: bool, remove_borders: bool, crop_variation: float, target_fps: float, output_dir: str): + video = cv2.VideoCapture(video_path) + vid_fps = self.__get_safe_fps(video, video_path) + max_length_frames = int(max_length * vid_fps) + min_length_frames = max(int(0.25 * vid_fps), 1) + total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) or 0 + timestamp_max_frame = min(self.__parse_timestamp_to_frames(timestamp_max, vid_fps), max(total_frames - 1, 0)) + timestamp_min_frame = min(self.__parse_timestamp_to_frames(timestamp_min, vid_fps), timestamp_max_frame) + + if split_at_cuts: + #use scenedetect to find cuts, based on start/end frame number + self.__update_status(f'Detecting scenes in "{os.path.basename(video_path)}"') + timecode_list = scenedetect.detect( + video_path=str(video_path), + detector=scenedetect.AdaptiveDetector(), + start_time=int(timestamp_min_frame), + end_time=int(timestamp_max_frame)) + scene_list = [(x[0].get_frames(), x[1].get_frames()) for x in timecode_list] + if not scene_list: + scene_list = [(timestamp_min_frame, timestamp_max_frame)] + else: + scene_list = [(timestamp_min_frame, timestamp_max_frame)] + + scene_list_split = [] + for scene in scene_list: + length = scene[1]-scene[0] + if length > max_length_frames: #check for any scenes longer than max length + n = math.ceil(length/max_length_frames) #divide into n new scenes + new_length = int(length/n) + new_splits = range(scene[0], scene[1]+min_length_frames, new_length) #divide clip into closest chunks to max_length + for i, _n in enumerate(new_splits[:-1]): + if new_splits[i + 1] - new_splits[i] > min_length_frames: + scene_list_split.append((new_splits[i], new_splits[i + 1])) + elif length > (min_length_frames + 2): + # Trim first/last frame to avoid transition artifacts + scene_list_split.append((scene[0] + 1, scene[1] - 1)) + + self.__update_status(f'Video "{os.path.basename(video_path)}" being split into {len(scene_list_split)} clips in "{output_dir}"') + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + futures = [ + executor.submit(self.__save_clip, scene, video_path, target_fps, + remove_borders, crop_variation, output_dir) + for scene in scene_list_split + ] + for future in concurrent.futures.as_completed(futures): + exc = future.exception() + if exc is not None: + self.__update_status(f'Error saving clip: {exc}') + + video.release() + + def __save_clip(self, scene: tuple[int, int], video_path: str, target_fps: float, + remove_borders: bool, crop_variation: float, output_dir: str): + basename, ext = os.path.splitext(os.path.basename(video_path)) + video = cv2.VideoCapture(str(video_path)) + fps = self.__get_safe_fps(video, video_path) + os.makedirs(output_dir, exist_ok=True) + output_name = f'{output_dir}{os.sep}{basename}_{scene[0]}-{scene[1]}' + output_ext = ".mp4" + + video.set(cv2.CAP_PROP_POS_FRAMES, (scene[1] + scene[0])//2) #set to middle of scene + frame_number = int(video.get(cv2.CAP_PROP_POS_FRAMES)) + success, frame = video.read() + if not success or frame is None: + self.__update_status(f'Failed to read frame from "{os.path.basename(video_path)}" at {int(frame_number)}. Skipping clip.') + video.release() + return + + # Blend random frames to detect borders, avoiding incorrect crop from black frames + if remove_borders: + frame_blend = frame + for i in range(5): + random_frame = random.randint(scene[0], scene[1]) + video.set(cv2.CAP_PROP_POS_FRAMES, random_frame) + success, frame = video.read() + if not success or frame is None: + continue + a = 1/(i+1) + b = 1-a + frame_blend = cv2.addWeighted(frame, a, frame_blend, b, 0) + x1, y1, w1, h1 = self.find_main_contour(frame_blend) + else: + x1 = 0 + y1 = 0 + h1, w1, _ = frame.shape + + y2, h2, x2, w2 = self.__get_random_aspect(h1, w1, crop_variation) + # Ensure dimensions are even, required + h2 -= h2 % 2 + w2 -= w2 % 2 + print(end='\x1b[2K') #clear terminal so next line can overwrite it + print(f'Saving frames {scene[0]}-{scene[1]} at size {w2}x{h2}', end="\r") + video.set(cv2.CAP_PROP_POS_FRAMES, (scene[1] + scene[0])//2) + success, frame = video.read() + if success: + try: + preview = Image.fromarray( + cv2.cvtColor(frame[y1+y2:y1+y2+h2, x1+x2:x1+x2+w2], cv2.COLOR_BGR2RGB)) + preview.thumbnail((150, 150)) + self.preview_image.configure(light_image=preview, size=preview.size) + #truncate filename of long files so UI doesn't shift around + filename_truncated = basename + ext if len(basename) < 20 else basename[:18] + ".." + ext + self.preview_image_label.configure( + text=f'{filename_truncated}\nFrames: {scene[0]}-{scene[1]}\nSize: {w2}x{h2}') + except Exception: + pass + video.release() + + if target_fps <= 0: + target_fps = fps + + output_path = f'{output_name}{output_ext}' + self.__write_clip_av(video_path, output_path, scene, fps, target_fps, + x1 + x2, y1 + y2, w2, h2) + + @staticmethod + def __write_clip_av(video_path: str, output_path: str, scene: tuple[int, int], + src_fps: float, target_fps: float, + crop_x: int, crop_y: int, crop_w: int, crop_h: int): + start_sec = scene[0] / src_fps + end_sec = scene[1] / src_fps + rate_frac = Fraction(target_fps).limit_denominator(10000) + stream_time_base = Fraction(rate_frac.denominator, rate_frac.numerator) + + with av.open(video_path) as input_container: + in_video = input_container.streams.video[0] + in_video.thread_type = 'AUTO' + in_audio = input_container.streams.audio[0] if input_container.streams.audio else None + + with av.open(output_path, mode='w') as output_container: + out_video = output_container.add_stream('libx264', rate=rate_frac) + out_video.width = crop_w + out_video.height = crop_h + out_video.pix_fmt = 'yuv420p' + out_video.time_base = stream_time_base + + out_audio = output_container.add_stream_from_template(in_audio) if in_audio else None + + input_container.seek(int(start_sec * 1_000_000)) + + out_frame_idx = 0 + out_time_step = 1.0 / target_fps + video_done = False + decode_streams = [s for s in (in_video, in_audio) if s is not None] + + for packet in input_container.demux(decode_streams): + if packet.stream == in_video: + if video_done: + continue + for frame in packet.decode(): + if frame.time is None or frame.time < start_sec: + continue + if frame.time >= end_sec: + video_done = True + break + + # FPS conversion: skip frames when source fps > target fps + if frame.time < start_sec + out_frame_idx * out_time_step: + continue + + img = frame.to_ndarray(format='bgr24') + cropped = img[crop_y:crop_y + crop_h, crop_x:crop_x + crop_w] + out_frame = av.VideoFrame.from_ndarray(cropped, format='bgr24') + out_frame.pts = out_frame_idx + out_frame.time_base = stream_time_base + + for out_pkt in out_video.encode(out_frame): + output_container.mux(out_pkt) + out_frame_idx += 1 + + elif packet.stream == in_audio and out_audio is not None: + if packet.dts is None: + continue + pkt_time = float(packet.pts * packet.time_base) + if pkt_time < start_sec or pkt_time >= end_sec: + continue + # Re-timestamp audio relative to clip start + packet.pts = int((pkt_time - start_sec) / packet.time_base) + packet.dts = packet.pts + packet.stream = out_audio + output_container.mux(packet) + + # Flush video encoder + for pkt in out_video.encode(): + output_container.mux(pkt) + + def __extract_images_button(self, batch_mode: bool): + self.__run_in_thread(self.__extract_images_multi, batch_mode) + + def __extract_images_multi(self, batch_mode : bool): + if not pathlib.Path(self.image_output_entry.get()).is_dir() or self.image_output_entry.get() == "": + self.__update_status("Invalid output directory!") + return + + # validate numeric inputs + try: + capture_rate = float(self.capture_rate_entry.get()) + blur_threshold = float(self.blur_threshold_entry.get()) + crop_variation = float(self.image_crop_entry.get()) + input_single_entry = self.image_single_entry.get() + input_multiple_entry = self.image_list_entry.get() + output_entry = self.image_output_entry.get() + except ValueError: + self.__update_status("Invalid numeric input for Images/sec, Blur Removal, or Crop Variation.") + return + if capture_rate <= 0: + self.__update_status("Images/sec must be > 0.") + return + if not (0.0 <= blur_threshold < 1.0): + self.__update_status("Blur Removal must be between 0.0 and 1.0.") + return + if not (0.0 <= crop_variation < 1.0): + self.__update_status("Crop Variation must be between 0.0 and 1.0.") + return + + input_videos = self.__get_vid_paths(batch_mode, input_single_entry, input_multiple_entry) + if not input_videos: + return + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + for video_path in input_videos: + output_directory = self.__get_output_dir( + self.output_subdir_img_entry.get(), batch_mode, + output_entry, video_path, input_multiple_entry) + time_start = "00:00:00" if batch_mode else str(self.image_time_start_entry.get()) + time_end = "99:99:99" if batch_mode else str(self.image_time_end_entry.get()) + executor.submit(self.__save_frames, + str(video_path), time_start, time_end, capture_rate, + blur_threshold, self.image_bordercrop.get(), + crop_variation, output_directory) + if batch_mode: + self.__update_status(f'Image extraction from all videos in {input_multiple_entry} complete') + else: + self.__update_status(f'Image extraction from "{input_single_entry}" complete') + + def __save_frames(self, video_path: str, timestamp_min: str, timestamp_max: str, capture_rate: float, + blur_threshold: float, remove_borders: bool, crop_variation: float, output_dir: str): + video = cv2.VideoCapture(video_path) + vid_fps = self.__get_safe_fps(video, video_path) + if capture_rate <= 0: + self.__update_status("Images/sec must be > 0.") + video.release() + return + image_rate = max(int(vid_fps / capture_rate), 1) + total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) or 0 + timestamp_max_frame = min(self.__parse_timestamp_to_frames(timestamp_max, vid_fps), max(total_frames - 1, 0)) + timestamp_min_frame = min(self.__parse_timestamp_to_frames(timestamp_min, vid_fps), timestamp_max_frame) + frame_range = range(timestamp_min_frame, timestamp_max_frame, image_rate) + frame_list = [] + + for n in frame_range: + #pick frame from random triangular distribution around center of each "chunk" of the video + frame = abs(int(random.triangular(n-(image_rate/2), n+(image_rate/2)))) + frame = max(0, min(frame, max(total_frames - 1, 0))) + frame_list.append(frame) + + self.__update_status(f'Video "{os.path.basename(video_path)}" will be split into {len(frame_list)} images in "{output_dir}"') + + output_list = [] + for f in frame_list: + video.set(cv2.CAP_PROP_POS_FRAMES, f) + success, frame = video.read() + if success and frame is not None: + frame_grayscale = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) + frame_sharpness = cv2.Laplacian(frame_grayscale, cv2.CV_64F).var() + output_list.append((f, frame_sharpness)) + + if not output_list: + self.__update_status(f'No frames extracted from "{os.path.basename(video_path)}" in the selected range.') + video.release() + return + + output_list_sorted = sorted(output_list, key=lambda x: x[1]) + cutoff = int(blur_threshold * len(output_list_sorted)) + output_list_cut = output_list_sorted[cutoff:] + self.__update_status(f'{cutoff} blurriest images have been dropped from "{os.path.basename(video_path)}"') + + basename, ext = os.path.splitext(os.path.basename(video_path)) + os.makedirs(output_dir, exist_ok=True) + + for f in output_list_cut: + filename = f'{output_dir}{os.sep}{basename}_{f[0]}.jpg' + video.set(cv2.CAP_PROP_POS_FRAMES, f[0]) + success, frame = video.read() + + #crop out borders of frame + if remove_borders and success and frame is not None: + x1, y1, w1, h1 = self.find_main_contour(frame) + frame_cropped = frame[y1:y1+h1, x1:x1+w1] + else: + frame_cropped = frame if success and frame is not None else None + if frame_cropped is not None: + x1 = 0 + y1 = 0 + h1, w1, _ = frame_cropped.shape + + y2, h2, x2, w2 = self.__get_random_aspect(h1, w1, crop_variation) + + if success and frame is not None and frame_cropped is not None: + print(end='\x1b[2K') #clear terminal so next line can overwrite it + print(f'Saving frame {f[0]} at size {w2}x{h2}', end="\r") + try: + preview = Image.fromarray( + cv2.cvtColor(frame_cropped[y2:y2+h2, x2:x2+w2], cv2.COLOR_BGR2RGB)) + preview.thumbnail((150, 150)) + filename_truncated = basename + ext if len(basename) < 20 else basename[:17] + "..." + ext + self.preview_image.configure(light_image=preview, size=preview.size) + self.preview_image_label.configure(text=f'{filename_truncated}\nFrame: {f[0]}\nSize: {w2}x{h2}') + except Exception: + pass # preview update is non-critical + + cv2.imwrite(filename, frame_cropped[y2:y2+h2, x2:x2+w2]) + video.release() + + def __download_button(self, batch_mode: bool): + self.__run_in_thread(self.__download_multi, batch_mode) + + def __update_status(self, status_text: str): + print(status_text) + self.status_label.configure(state="normal") + self.status_label.insert(index="end", text=status_text + "\n") + self.status_label.configure(state="disabled") + + def __download_multi(self, batch_mode: bool): + if not pathlib.Path(self.download_output_entry.get()).is_dir() or self.download_output_entry.get() == "": + self.__update_status("Invalid output directory!") + return + + if not batch_mode: + ydl_urls = [self.download_link_entry.get()] + elif batch_mode: + ydl_path = pathlib.Path(self.download_list_entry.get()) + if ydl_path.is_file() and ydl_path.suffix.lower() == ".txt": + with open(ydl_path) as file: + ydl_urls = file.readlines() + else: + self.__update_status("Invalid link list!") + return + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + for url in ydl_urls: + executor.submit(self.__download_video, + url.strip(), self.download_output_entry.get(), + self.download_args_entry.get("0.0", ctk.END)) + + self.__update_status(f'Completed {len(ydl_urls)} downloads.') + + def __download_video(self, url: str, output_dir: str, output_args: str): + url = (url or "").strip() + if not url: + self.__update_status("Empty URL, skipping download.") + return + + #Respect quotes and split into list to run as yt-dlp command + additional_args = shlex.split(output_args.strip()) if output_args and output_args.strip() else [] + cmd = ["yt-dlp", "-o", "%(title)s.%(ext)s", "-P", output_dir] + additional_args + [url] + + self.__update_status(f'Downloading {url}') + subprocess.run(cmd) + self.__update_status(f'Download {url} done!') diff --git a/modules/ui/GenerateCaptionsWindowController.py b/modules/ui/GenerateCaptionsWindowController.py new file mode 100644 index 000000000..1690879f1 --- /dev/null +++ b/modules/ui/GenerateCaptionsWindowController.py @@ -0,0 +1,133 @@ +import contextlib +import tkinter as tk +from tkinter import filedialog + +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk + + +class GenerateCaptionsWindow(ctk.CTkToplevel): + def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs): + """ + Window for generating captions for a folder of images + + Parameters: + parent (`Tk`): the parent window + path (`str`): the path to the folder + parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox + """ + super().__init__(parent, *args, **kwargs) + self.parent = parent + + if path is None: + path = "" + + self.mode_var = ctk.StringVar(self, "Create if absent") + self.modes = ["Replace all captions", "Create if absent", "Add as new line"] + self.model_var = ctk.StringVar(self, "Blip") + self.models = ["Blip", "Blip2", "WD14 VIT v2"] + + self.title("Batch generate captions") + self.geometry("360x360") + self.resizable(True, True) + + self.frame = ctk.CTkFrame(self, width=600, height=300) + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + + self.model_label = ctk.CTkLabel(self.frame, text="Model", width=100) + self.model_label.grid(row=0, column=0, sticky="w", padx=5, pady=5) + self.model_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.model_var, values=self.models, dynamic_resizing=False, width=200) + self.model_dropdown.grid(row=0, column=1, sticky="w", padx=5, pady=5) + + self.path_label = ctk.CTkLabel(self.frame, text="Folder", width=100) + self.path_label.grid(row=1, column=0, sticky="w",padx=5, pady=5) + self.path_entry = ctk.CTkEntry(self.frame, width=150) + self.path_entry.insert(0, path) + self.path_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.path_button = ctk.CTkButton(self.frame, width=30, text="...", command=lambda: self.browse_for_path(self.path_entry)) + self.path_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) + + self.caption_label = ctk.CTkLabel(self.frame, text="Initial Caption", width=100) + self.caption_label.grid(row=2, column=0, sticky="w", padx=5, pady=5) + self.caption_entry = ctk.CTkEntry(self.frame, width=200) + self.caption_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + + self.prefix_label = ctk.CTkLabel(self.frame, text="Caption Prefix", width=100) + self.prefix_label.grid(row=3, column=0, sticky="w", padx=5, pady=5) + self.prefix_entry = ctk.CTkEntry(self.frame, width=200) + self.prefix_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) + + self.postfix_label = ctk.CTkLabel(self.frame, text="Caption Postfix", width=100) + self.postfix_label.grid(row=4, column=0, sticky="w", padx=5, pady=5) + self.postfix_entry = ctk.CTkEntry(self.frame, width=200) + self.postfix_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + + self.mode_label = ctk.CTkLabel(self.frame, text="Mode", width=100) + self.mode_label.grid(row=5, column=0, sticky="w", padx=5, pady=5) + self.mode_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.mode_var, values=self.modes, dynamic_resizing=False, width=200) + self.mode_dropdown.grid(row=5, column=1, sticky="w", padx=5, pady=5) + + self.include_subdirectories_label = ctk.CTkLabel(self.frame, text="Include subfolders", width=100) + self.include_subdirectories_label.grid(row=6, column=0, sticky="w", padx=5, pady=5) + self.include_subdirectories_var = ctk.BooleanVar(self, parent_include_subdirectories) + self.include_subdirectories_switch = ctk.CTkSwitch(self.frame, text="", variable=self.include_subdirectories_var) + self.include_subdirectories_switch.grid(row=6, column=1, sticky="w", padx=5, pady=5) + + self.progress_label = ctk.CTkLabel(self.frame, text="Progress: 0/0", width=100) + self.progress_label.grid(row=7, column=0, sticky="w", padx=5, pady=5) + self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) + self.progress.grid(row=7, column=1, sticky="w", padx=5, pady=5) + + self.create_captions_button = ctk.CTkButton(self.frame, text="Create Captions", width=310, command=self.create_captions) + self.create_captions_button.grid(row=8, column=0, columnspan=2, sticky="w", padx=5, pady=5) + + self.frame.pack(fill="both", expand=True) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def browse_for_path(self, entry_box): + # get the path from the user + path = filedialog.askdirectory() + # set the path to the entry box + # delete entry box text + entry_box.focus_set() + entry_box.delete(0, filedialog.END) + entry_box.insert(0, path) + self.focus_set() + + def set_progress(self, value, max_value): + progress = value / max_value + self.progress.set(progress) + self.progress_label.configure(text=f"{value}/{max_value}") + self.progress.update() + + def create_captions(self): + self.parent.load_captioning_model(self.model_var.get()) + + mode = { + "Replace all captions": "replace", + "Create if absent": "fill", + "Add as new line": "add", + }[self.mode_var.get()] + + self.parent.captioning_model.caption_folder( + sample_dir=self.path_entry.get(), + initial_caption=self.caption_entry.get(), + caption_prefix=self.prefix_entry.get(), + caption_postfix=self.postfix_entry.get(), + mode=mode, + progress_callback=self.set_progress, + include_subdirectories=self.include_subdirectories_var.get(), + ) + self.parent.load_image() + + def destroy(self): + with contextlib.suppress(tk.TclError): + self.grab_release() + + super().destroy() diff --git a/modules/ui/GenerateMasksWindowController.py b/modules/ui/GenerateMasksWindowController.py new file mode 100644 index 000000000..daff0d3d5 --- /dev/null +++ b/modules/ui/GenerateMasksWindowController.py @@ -0,0 +1,151 @@ +import contextlib +import tkinter as tk +from tkinter import filedialog + +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk + + +class GenerateMasksWindow(ctk.CTkToplevel): + def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs): + """ + Window for generating masks for a folder of images + + Parameters: + parent (`Tk`): the parent window + path (`str`): the path to the folder + parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox + """ + super().__init__(parent, *args, **kwargs) + + self.parent = parent + if path is None: + path = "" + + self.mode_var = ctk.StringVar(self, "Create if absent") + self.modes = ["Replace all masks", "Create if absent", "Add to existing", "Subtract from existing", "Blend with existing"] + self.model_var = ctk.StringVar(self, "ClipSeg") + self.models = ["ClipSeg", "Rembg", "Rembg-Human", "Hex Color"] + + self.title("Batch generate masks") + self.geometry("360x430") + self.resizable(True, True) + + self.frame = ctk.CTkFrame(self, width=600, height=300) + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + + self.model_label = ctk.CTkLabel(self.frame, text="Model", width=100) + self.model_label.grid(row=0, column=0, sticky="w", padx=5, pady=5) + self.model_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.model_var, values=self.models, dynamic_resizing=False, width=200) + self.model_dropdown.grid(row=0, column=1, sticky="w", padx=5, pady=5) + + self.path_label = ctk.CTkLabel(self.frame, text="Folder", width=100) + self.path_label.grid(row=1, column=0, sticky="w",padx=5, pady=5) + self.path_entry = ctk.CTkEntry(self.frame, width=150) + self.path_entry.insert(0, path) + self.path_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.path_button = ctk.CTkButton(self.frame, width=30, text="...", command=lambda: self.browse_for_path(self.path_entry)) + self.path_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) + + self.prompt_label = ctk.CTkLabel(self.frame, text="Prompt", width=100) + self.prompt_label.grid(row=2, column=0, sticky="w",padx=5, pady=5) + self.prompt_entry = ctk.CTkEntry(self.frame, width=200) + self.prompt_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + + self.mode_label = ctk.CTkLabel(self.frame, text="Mode", width=100) + self.mode_label.grid(row=3, column=0, sticky="w", padx=5, pady=5) + self.mode_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.mode_var, values=self.modes, dynamic_resizing=False, width=200) + self.mode_dropdown.grid(row=3, column=1, sticky="w", padx=5, pady=5) + + self.threshold_label = ctk.CTkLabel(self.frame, text="Threshold", width=100) + self.threshold_label.grid(row=4, column=0, sticky="w", padx=5, pady=5) + self.threshold_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="0.0 - 1.0") + self.threshold_entry.insert(0, "0.3") + self.threshold_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + + self.smooth_label = ctk.CTkLabel(self.frame, text="Smooth", width=100) + self.smooth_label.grid(row=5, column=0, sticky="w", padx=5, pady=5) + self.smooth_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="5") + self.smooth_entry.insert(0, 5) + self.smooth_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) + + self.expand_label = ctk.CTkLabel(self.frame, text="Expand", width=100) + self.expand_label.grid(row=6, column=0, sticky="w", padx=5, pady=5) + self.expand_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="10") + self.expand_entry.insert(0, 10) + self.expand_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) + + self.alpha_label = ctk.CTkLabel(self.frame, text="Alpha", width=100) + self.alpha_label.grid(row=7, column=0, sticky="w", padx=5, pady=5) + self.alpha_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="1") + self.alpha_entry.insert(0, 1) + self.alpha_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) + + self.include_subdirectories_label = ctk.CTkLabel(self.frame, text="Include subfolders", width=100) + self.include_subdirectories_label.grid(row=8, column=0, sticky="w", padx=5, pady=5) + self.include_subdirectories_var = ctk.BooleanVar(self, parent_include_subdirectories) + self.include_subdirectories_switch = ctk.CTkSwitch(self.frame, text="", variable=self.include_subdirectories_var) + self.include_subdirectories_switch.grid(row=8, column=1, sticky="w", padx=5, pady=5) + + self.progress_label = ctk.CTkLabel(self.frame, text="Progress: 0/0", width=100) + self.progress_label.grid(row=9, column=0, sticky="w", padx=5, pady=5) + self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) + self.progress.grid(row=9, column=1, sticky="w", padx=5, pady=5) + + self.create_masks_button = ctk.CTkButton(self.frame, text="Create Masks", width=310, command=self.create_masks) + self.create_masks_button.grid(row=10, column=0, columnspan=2, sticky="w", padx=5, pady=5) + + self.frame.pack(fill="both", expand=True) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def browse_for_path(self, entry_box): + # get the path from the user + path = filedialog.askdirectory() + # set the path to the entry box + # delete entry box text + entry_box.focus_set() + entry_box.delete(0, filedialog.END) + entry_box.insert(0, path) + self.focus_set() + + def set_progress(self, value, max_value): + progress = value / max_value + self.progress.set(progress) + self.progress_label.configure(text=f"{value}/{max_value}") + self.progress.update() + + def create_masks(self): + self.parent.load_masking_model(self.model_var.get()) + + mode = { + "Replace all masks": "replace", + "Create if absent": "fill", + "Add to existing": "add", + "Subtract from existing": "subtract", + "Blend with existing": "blend", + }[self.mode_var.get()] + + self.parent.masking_model.mask_folder( + sample_dir=self.path_entry.get(), + prompts=[self.prompt_entry.get()], + mode=mode, + alpha=float(self.alpha_entry.get()), + threshold=float(self.threshold_entry.get()), + smooth_pixels=int(self.smooth_entry.get()), + expand_pixels=int(self.expand_entry.get()), + progress_callback=self.set_progress, + include_subdirectories=self.include_subdirectories_var.get(), + ) + self.parent.load_image() + + def destroy(self): + with contextlib.suppress(tk.TclError): + self.grab_release() + + super().destroy() diff --git a/modules/ui/OptimizerParamsWindowController.py b/modules/ui/OptimizerParamsWindowController.py new file mode 100644 index 000000000..16063c26c --- /dev/null +++ b/modules/ui/OptimizerParamsWindowController.py @@ -0,0 +1,288 @@ +import contextlib +from tkinter import TclError + +from modules.ui.MuonAdamWindow import MUON_AUX_ADAM_DEFAULTS, MuonAdamWindow +from modules.util.config.TrainConfig import TrainConfig, TrainOptimizerConfig +from modules.util.enum.Optimizer import Optimizer +from modules.util.optimizer_util import ( + OPTIMIZER_DEFAULT_PARAMETERS, + change_optimizer, + load_optimizer_defaults, + update_optimizer_config, +) +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class OptimizerParamsWindow(ctk.CTkToplevel): + def __init__( + self, + parent, + train_config: TrainConfig, + ui_state, + *args, **kwargs, + ): + super().__init__(parent, *args, **kwargs) + + self.parent = parent + self.train_config = train_config + self.ui_state = ui_state + self.optimizer_ui_state = ui_state.get_var("optimizer") + self.protocol("WM_DELETE_WINDOW", self.on_window_close) + self.muon_adam_button = None + + self.title("Optimizer Settings") + self.geometry("800x500") + self.resizable(True, True) + + self.grid_rowconfigure(0, weight=1) + self.grid_rowconfigure(1, weight=0) + self.grid_columnconfigure(0, weight=1) + + self.frame = ctk.CTkScrollableFrame(self, fg_color="transparent") + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + + self.frame.grid_columnconfigure(0, weight=0) + self.frame.grid_columnconfigure(1, weight=1) + self.frame.grid_columnconfigure(2, minsize=50) + self.frame.grid_columnconfigure(3, weight=0) + self.frame.grid_columnconfigure(4, weight=1) + + components.button(self, 1, 0, "ok", command=self.on_window_close) + self.main_frame(self.frame) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def main_frame(self, master): + # Optimizer + components.label(master, 0, 0, "Optimizer", + tooltip="The type of optimizer") + + # Create the optimizer dropdown menu and set the command + components.options(master, 0, 1, [str(x) for x in list(Optimizer)], self.optimizer_ui_state, "optimizer", + command=self.on_optimizer_change) + + # Defaults Button + components.label(master, 0, 3, "Optimizer Defaults", + tooltip="Load default settings for the selected optimizer") + components.button(self.frame, 0, 4, "Load Defaults", self.load_defaults, + tooltip="Load default settings for the selected optimizer") + + self.create_dynamic_ui(master) + + def clear_dynamic_ui(self, master): + with contextlib.suppress(TclError): + for widget in master.winfo_children(): + grid_info = widget.grid_info() + if int(grid_info["row"]) >= 1: + widget.destroy() + + def create_dynamic_ui( + self, + master, + ): + + # Lookup for the title and tooltip for a key + # @formatter:off + KEY_DETAIL_MAP = { + 'adam_w_mode': {'title': 'Adam W Mode', 'tooltip': 'Whether to use weight decay correction for Adam optimizer.', 'type': 'bool'}, + 'alpha': {'title': 'Alpha', 'tooltip': 'Smoothing parameter for RMSprop and others.', 'type': 'float'}, + 'amsgrad': {'title': 'AMSGrad', 'tooltip': 'Whether to use the AMSGrad variant for Adam.', 'type': 'bool'}, + 'beta1': {'title': 'Beta1', 'tooltip': 'optimizer_momentum term.', 'type': 'float'}, + 'beta2': {'title': 'Beta2', 'tooltip': 'Coefficients for computing running averages of gradient.', 'type': 'float'}, + 'beta3': {'title': 'Beta3', 'tooltip': 'Coefficient for computing the Prodigy stepsize.', 'type': 'float'}, + 'bias_correction': {'title': 'Bias Correction', 'tooltip': 'Whether to use bias correction in optimization algorithms like Adam.', 'type': 'bool'}, + 'block_wise': {'title': 'Block Wise', 'tooltip': 'Whether to perform block-wise model update.', 'type': 'bool'}, + 'capturable': {'title': 'Capturable', 'tooltip': 'Whether some property of the optimizer can be captured.', 'type': 'bool'}, + 'centered': {'title': 'Centered', 'tooltip': 'Whether to center the gradient before scaling. Great for stabilizing the training process.', 'type': 'bool'}, + 'clip_threshold': {'title': 'Clip Threshold', 'tooltip': 'Clipping value for gradients.', 'type': 'float'}, + 'd0': {'title': 'Initial D', 'tooltip': 'Initial D estimate for D-adaptation.', 'type': 'float'}, + 'd_coef': {'title': 'D Coefficient', 'tooltip': 'Coefficient in the expression for the estimate of d.', 'type': 'float'}, + 'dampening': {'title': 'Dampening', 'tooltip': 'Dampening for optimizer_momentum.', 'type': 'float'}, + 'decay_rate': {'title': 'Decay Rate', 'tooltip': 'Rate of decay for moment estimation.', 'type': 'float'}, + 'decouple': {'title': 'Decouple', 'tooltip': 'Use AdamW style optimizer_decoupled weight decay.', 'type': 'bool'}, + 'differentiable': {'title': 'Differentiable', 'tooltip': 'Whether the optimization function is optimizer_differentiable.', 'type': 'bool'}, + 'eps': {'title': 'EPS', 'tooltip': 'A small value to prevent division by zero.', 'type': 'float'}, + 'eps2': {'title': 'EPS 2', 'tooltip': 'A small value to prevent division by zero.', 'type': 'float'}, + 'foreach': {'title': 'ForEach', 'tooltip': 'Whether to use a foreach implementation if available. This implementation is usually faster.', 'type': 'bool'}, + 'fsdp_in_use': {'title': 'FSDP in Use', 'tooltip': 'Flag for using sharded parameters.', 'type': 'bool'}, + 'fused': {'title': 'Fused', 'tooltip': 'Whether to use a fused implementation if available. This implementation is usually faster and requires less memory.', 'type': 'bool'}, + 'fused_back_pass': {'title': 'Fused Back Pass', 'tooltip': 'Whether to fuse the back propagation pass with the optimizer step. This reduces VRAM usage, but is not compatible with gradient accumulation.', 'type': 'bool'}, + 'growth_rate': {'title': 'Growth Rate', 'tooltip': 'Limit for D estimate growth rate.', 'type': 'float'}, + 'initial_accumulator_value': {'title': 'Initial Accumulator Value', 'tooltip': 'Initial value for Adagrad optimizer.', 'type': 'float'}, + 'initial_accumulator': {'title': 'Initial Accumulator', 'tooltip': 'Sets the starting value for both moment estimates to ensure numerical stability and balanced adaptive updates early in training.', 'type': 'float'}, + 'is_paged': {'title': 'Is Paged', 'tooltip': 'Whether the optimizer\'s internal state should be paged to CPU.', 'type': 'bool'}, + 'log_every': {'title': 'Log Every', 'tooltip': 'Intervals at which logging should occur.', 'type': 'int'}, + 'lr_decay': {'title': 'LR Decay', 'tooltip': 'Rate at which learning rate decreases.', 'type': 'float'}, + 'max_unorm': {'title': 'Max Unorm', 'tooltip': 'Maximum value for gradient clipping by norms.', 'type': 'float'}, + 'maximize': {'title': 'Maximize', 'tooltip': 'Whether to optimizer_maximize the optimization function.', 'type': 'bool'}, + 'min_8bit_size': {'title': 'Min 8bit Size', 'tooltip': 'Minimum tensor size for 8-bit quantization.', 'type': 'int'}, + 'quant_block_size': {'title': 'Quant Block Size', 'tooltip': 'Size of a block of normalized 8-bit quantization data. Larger values increase memory efficiency at the cost of data precision.', 'type': 'int'}, + 'momentum': {'title': 'optimizer_momentum', 'tooltip': 'Factor to accelerate SGD in relevant direction.', 'type': 'float'}, + 'nesterov': {'title': 'Nesterov', 'tooltip': 'Whether to enable Nesterov optimizer_momentum.', 'type': 'bool'}, + 'no_prox': {'title': 'No Prox', 'tooltip': 'Whether to use proximity updates or not.', 'type': 'bool'}, + 'optim_bits': {'title': 'Optim Bits', 'tooltip': 'Number of bits used for optimization.', 'type': 'int'}, + 'percentile_clipping': {'title': 'Percentile Clipping', 'tooltip': 'Gradient clipping based on percentile values.', 'type': 'int'}, + 'relative_step': {'title': 'Relative Step', 'tooltip': 'Whether to use a relative step size.', 'type': 'bool'}, + 'safeguard_warmup': {'title': 'Safeguard Warmup', 'tooltip': 'Avoid issues during warm-up stage.', 'type': 'bool'}, + 'scale_parameter': {'title': 'Scale Parameter', 'tooltip': 'Whether to scale the parameter or not.', 'type': 'bool'}, + 'stochastic_rounding': {'title': 'Stochastic Rounding', 'tooltip': 'Stochastic rounding for weight updates. Improves quality when using bfloat16 weights.', 'type': 'bool'}, + 'use_bias_correction': {'title': 'Bias Correction', 'tooltip': 'Turn on Adam\'s bias correction.', 'type': 'bool'}, + 'use_triton': {'title': 'Use Triton', 'tooltip': 'Whether Triton optimization should be used.', 'type': 'bool'}, + 'warmup_init': {'title': 'Warmup Initialization', 'tooltip': 'Whether to warm-up the optimizer initialization.', 'type': 'bool'}, + 'weight_decay': {'title': 'Weight Decay', 'tooltip': 'Regularization to prevent overfitting.', 'type': 'float'}, + 'weight_lr_power': {'title': 'Weight LR Power', 'tooltip': 'During warmup, the weights in the average will be equal to lr raised to this power. Set to 0 for no weighting.', 'type': 'float'}, + 'decoupled_decay': {'title': 'Decoupled Decay', 'tooltip': 'If set as True, then the optimizer uses decoupled weight decay as in AdamW.', 'type': 'bool'}, + 'fixed_decay': {'title': 'Fixed Decay', 'tooltip': '(When Decoupled Decay is True:) Applies fixed weight decay when True; scales decay with learning rate when False.', 'type': 'bool'}, + 'rectify': {'title': 'Rectify', 'tooltip': 'Perform the rectified update similar to RAdam.', 'type': 'bool'}, + 'degenerated_to_sgd': {'title': 'Degenerated to SGD', 'tooltip': 'Performs SGD update when gradient variance is high.', 'type': 'bool'}, + 'k': {'title': 'K', 'tooltip': 'Number of vector projected per iteration.', 'type': 'int'}, + 'xi': {'title': 'Xi', 'tooltip': 'Term used in vector projections to avoid division by zero.', 'type': 'float'}, + 'n_sma_threshold': {'title': 'N SMA Threshold', 'tooltip': 'Number of SMA threshold.', 'type': 'int'}, + 'ams_bound': {'title': 'AMS Bound', 'tooltip': 'Whether to use the AMSBound variant.', 'type': 'bool'}, + 'r': {'title': 'R', 'tooltip': 'EMA factor.', 'type': 'float'}, + 'adanorm': {'title': 'AdaNorm', 'tooltip': 'Whether to use the AdaNorm variant', 'type': 'bool'}, + 'adam_debias': {'title': 'Adam Debias', 'tooltip': 'Only correct the denominator to avoid inflating step sizes early in training.', 'type': 'bool'}, + 'slice_p': {'title': 'Slice parameters', 'tooltip': 'Reduce memory usage by calculating LR adaptation statistics on only every pth entry of each tensor. For values greater than 1 this is an approximation to standard Prodigy. Values ~11 are reasonable.', 'type': 'int'}, + 'cautious': {'title': 'Cautious', 'tooltip': 'Whether to use the Cautious variant', 'type': 'bool'}, + 'weight_decay_by_lr': {'title': 'weight_decay_by_lr', 'tooltip': 'Automatically adjust weight decay based on lr', 'type': 'bool'}, + 'prodigy_steps': {'title': 'prodigy_steps', 'tooltip': 'Turn off Prodigy after N steps', 'type': 'int'}, + 'use_speed': {'title': 'use_speed', 'tooltip': 'use_speed method', 'type': 'bool'}, + 'split_groups': {'title': 'split_groups', 'tooltip': 'Use split groups when training multiple params(uNet,TE..)', 'type': 'bool'}, + 'split_groups_mean': {'title': 'split_groups_mean', 'tooltip': 'Use mean for split groups', 'type': 'bool'}, + 'factored': {'title': 'factored', 'tooltip': 'Use factored', 'type': 'bool'}, + 'factored_fp32': {'title': 'factored_fp32', 'tooltip': 'Use factored_fp32', 'type': 'bool'}, + 'use_stableadamw': {'title': 'use_stableadamw', 'tooltip': 'Use use_stableadamw for gradient scaling', 'type': 'bool'}, + 'use_cautious': {'title': 'use_cautious', 'tooltip': 'Use cautious method', 'type': 'bool'}, + 'use_grams': {'title': 'use_grams', 'tooltip': 'Use grams method', 'type': 'bool'}, + 'use_adopt': {'title': 'use_adopt', 'tooltip': 'Use adopt method', 'type': 'bool'}, + 'd_limiter': {'title': 'd_limiter', 'tooltip': 'Prevent over-estimated LRs when gradients and EMA are still stabilizing', 'type': 'bool'}, + 'use_schedulefree': {'title': 'use_schedulefree', 'tooltip': 'Use Schedulefree method', 'type': 'bool'}, + 'use_orthograd': {'title': 'use_orthograd', 'tooltip': 'Use orthograd method', 'type': 'bool'}, + 'nnmf_factor': {'title': 'Factored Optimizer', 'tooltip': 'Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. This is an approximation which may slightly alter training dynamics.', 'type': 'bool'}, + 'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization.', 'type': 'bool'}, + 'use_atan2': {'title': 'Atan2 Scaling', 'tooltip': 'A robust replacement for eps, which also incorporates gradient clipping, bounding and stabilizing the optimizer updates.', 'type': 'bool'}, + 'use_AdEMAMix': {'title': 'AdEMAMix EMA', 'tooltip': 'Adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, and accelerate the training.', 'type': 'bool'}, + 'beta3_ema': {'title': 'Beta3 EMA', 'tooltip': 'Coefficient for slow-moving EMA of AdEMAMix.', 'type': 'float'}, + 'beta1_warmup': {'title': 'Beta1 Warmup Steps', 'tooltip': 'Number of warmup steps to gradually increase beta1 from Minimum Beta1 Value to its final value. During warmup, beta1 increases linearly. leave it empty to disable warmup and use constant beta1.', 'type': 'int'}, + 'min_beta1': {'title': 'Minimum Beta1', 'tooltip': 'Starting beta1 value for warmup scheduling. Used only when beta1 warmup is enabled. Lower values allow faster initial adaptation, while higher values provide more smoothing. The final beta1 value is specified in the beta1 parameter.', 'type': 'float'}, + 'Simplified_AdEMAMix': {'title': 'Simplified AdEMAMix', 'tooltip': "Enables a simplified, single-EMA variant of AdEMAMix. Instead of blending two moving averages (fast and slow momentum), this version combines the raw current gradient (controlled by 'Grad α') directly with a single theory-based momentum. This makes the optimizer highly responsive to recent gradient information, which can accelerate training in all batch size scenarios when tuned correctly.", 'type': 'bool'}, + 'alpha_grad': {'title': 'Grad α', 'tooltip': 'Controls the mixing coefficient between raw gradients and momentum gradients in Simplified AdEMAMix. Higher values (e.g., 10-100) emphasize recent gradients, suitable for small batch sizes to reduce noise. Lower values (e.g., 0-1) emphasize historical gradients, suitable for large batch sizes for stability. Setting to 0 uses only momentum gradients without raw gradient contribution.', 'type': 'float'}, + 'kourkoutas_beta': {'title': 'Kourkoutas Beta', 'tooltip': 'Enables a layer-wise dynamic β₂ adaptation. This feature makes the optimizer more responsive to "spiky" gradients by lowering β₂ during periods of high variance, and more stable during calm periods by raising β₂ towards its maximum. It can significantly improve training stability and final loss.', 'type': 'bool'}, + 'schedulefree_c': {'title': 'Schedule free averaging strength', 'tooltip': 'Larger values = more responsive (shorter averaging window); smaller values = smoother (longer window). Set to 0 to disable and use the original Schedule-Free rule. Short small batches (≈6-12); long/large-batch (≈50-200).', 'type': 'float'}, + 'ns_steps': {'title': 'Newton-Schulz Iterations', 'tooltip': 'Controls the number of iterations for update orthogonalization. Higher values improve the updates quality but make each step slower. Lower values are faster per step but may be less effective.', 'type': 'int'}, + 'MuonWithAuxAdam': {'title': 'MuonWithAuxAdam', 'tooltip': 'Whether to use the standard way of Muon. Non-hidden layers fallback to ADAMW, and MUON takes the rest. Note: The auxiliary Adam (ADAMW) is typically only relevant for training "full" LoRA (LoRA for all layers) or full finetune and is irrelevant for most common LoRA use cases.', 'type': 'bool'}, + 'muon_hidden_layers': {'title': 'Hidden Layers', 'tooltip': 'Comma-separated list of hidden layers to train using Muon. Regular expressions (if toggled) are supported. Any model layer with a matching name will be trained using Muon. If None is provided it will default to using automatic way of finding hidden layers.', 'type': 'str'}, + 'muon_adam_regex': {'title': 'Use Regex', 'tooltip': 'Whether to use regular expressions for hidden layers.', 'type': 'bool'}, + 'muon_adam_lr': {'title': 'Auxiliary Adam LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer. If empty, it will use the main learning rate.', 'type': 'float'}, + 'muon_te1_adam_lr': {'title': 'AuxAdam TE1 LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer for the first text encoder. If empty, it will use the Auxiliary Adam LR.', 'type': 'float'}, + 'muon_te2_adam_lr': {'title': 'AuxAdam TE2 LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer for the second text encoder. If empty, it will use the Auxiliary Adam LR.', 'type': 'float'}, + 'rms_rescaling': {'title': 'RMS Rescaling', 'tooltip': 'Muon already scales its updates to approximate and use the same learning rate (LR) as Adam. This option integrates a more accurate method to match the Adam LR, but it is slower.', 'type': 'bool'}, + 'normuon_variant': {'title': 'NorMuon Variant', 'tooltip': 'Enables the NorMuon optimizer variant, which combines Muon orthogonalization with per-neuron adaptive learning rates for better convergence and balanced parameter updates. Costs only one scalar state buffer per parameter group, size few KBs, maintaining high memory efficiency.', 'type': 'bool'}, + 'beta2_normuon': {'title': 'NorMuon Beta2', 'tooltip': 'Exponential decay rate for the neuron-wise second-moment estimator in NorMuon (analogous to Adams beta2). Controls how past squared updates influence current normalization.', 'type': 'float'}, + 'low_rank_ortho': {'title': 'Low-rank Orthogonalization', 'tooltip': 'Use low-rank orthogonalization to accelerate Muon by orthogonalizing only in a low-dimensional subspace, improving speed and noise robustness.', 'type': 'bool'}, + 'ortho_rank': {'title': 'Ortho Rank', 'tooltip': 'Target rank for low-rank orthogonalization. Controls the dimensionality of the subspace used for efficient and noise-robust orthogonalization.', 'type': 'int'}, + 'accelerated_ns': {'title': 'Accelerated Newton-Schulz', 'tooltip': 'Applies an enhanced Newton-Schulz variant that replaces heuristic coefficients with optimal coefficients derived at each step. This improves performance and convergence by reducing the number of required operations.', 'type': 'bool'}, + 'cautious_wd': {'title': 'Cautious Weight Decay', 'tooltip': 'Applies weight decay only to parameter coordinates whose signs align with the optimizer update direction. This preserves the original optimization objective while still benefiting from regularization effects, leading to improved convergence and better final performance.', 'type': 'bool'}, + 'approx_mars': {'title': 'Approx MARS-M', 'tooltip': 'Enables Approximated MARS-M, a variance reduction technique. It uses the previous step\'s gradient to correct the current update, leading to lower losses and improved convergence stability. This requires additional state to store the previous gradient.', 'type': 'bool'}, + 'auto_kappa_p': {'title': 'Auto Lion-K', 'tooltip': 'Automatically determines the optimal P-value based on layer dimensions. Uses p=2.0 (Spherical) for 4D (Conv) tensors for stability and rotational invariance, and p=1.0 (Sign) for 2D (Linear) tensors for sparsity. Overrides the manual P-value. Recommend for unet models.', 'type': 'bool'}, + 'compile': {'title': 'Compiled Optimizer', 'tooltip': 'Enables PyTorch compilation for the optimizer internal step logic. This is intended to improve performance by allowing PyTorch to fuse operations and optimize the computational graph.', 'type': 'bool'}, + } + # @formatter:on + + if not self.winfo_exists(): # check if this window isn't open + return + + selected_optimizer = self.train_config.optimizer.optimizer + + # Extract the keys for the selected optimizer + for index, key in enumerate(OPTIMIZER_DEFAULT_PARAMETERS[selected_optimizer].keys()): + if key not in KEY_DETAIL_MAP: + continue + arg_info = KEY_DETAIL_MAP[key] + + title = arg_info['title'] + tooltip = arg_info['tooltip'] + type = arg_info['type'] + + row = (index // 2) + 1 + col = 3 * (index % 2) + + components.label(master, row, col, title, tooltip=tooltip) + + if key == 'MuonWithAuxAdam': + frame = ctk.CTkFrame(master, fg_color="transparent") + frame.grid(row=row, column=col + 1, columnspan=2, sticky="ew", padx=0, pady=0) + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=0) + + components.switch(frame, 0, 0, self.optimizer_ui_state, key, command=self.update_user_pref) + + self.muon_adam_button = components.button( + frame, 0, 1, "...", self.open_muon_adam_window, + tooltip="Configure the auxiliary AdamW_adv optimizer", + width=20, padx=5 ) + self.toggle_muon_adam_button() + elif type != 'bool': + components.entry(master, row, col + 1, self.optimizer_ui_state, key, + command=self.update_user_pref) + else: + components.switch(master, row, col + 1, self.optimizer_ui_state, key, + command=self.update_user_pref) + + def update_user_pref(self, *args): + update_optimizer_config(self.train_config) + self.toggle_muon_adam_button() + + def on_optimizer_change(self, *args): + optimizer_config = change_optimizer(self.train_config) + self.ui_state.get_var("optimizer").update(optimizer_config) + + self.clear_dynamic_ui(self.frame) + self.create_dynamic_ui(self.frame) + + def load_defaults(self, *args): + optimizer_config = load_optimizer_defaults(self.train_config) + self.ui_state.get_var("optimizer").update(optimizer_config) + + def on_window_close(self): + self.destroy() + + def toggle_muon_adam_button(self): + if self.muon_adam_button and self.muon_adam_button.winfo_exists(): + muon_with_adam = self.optimizer_ui_state.get_var("MuonWithAuxAdam").get() + self.muon_adam_button.configure(state="normal" if muon_with_adam else "disabled") + + def open_muon_adam_window(self): + current_optimizer = self.train_config.optimizer.optimizer + + adam_config = TrainOptimizerConfig.default_values() + current_state = self.train_config.optimizer.muon_adam_config + + if current_optimizer == Optimizer.MUON: + defaults = MUON_AUX_ADAM_DEFAULTS + else: + defaults = OPTIMIZER_DEFAULT_PARAMETERS[Optimizer.ADAMW_ADV] + + if current_state is None: + adam_config.from_dict(defaults) + if current_optimizer != Optimizer.MUON: + adam_config.optimizer = Optimizer.ADAMW_ADV + elif isinstance(current_state, dict): + adam_config.from_dict(current_state) + else: + # Should not happen if TrainConfig defines it as dict, but for safety + adam_config = current_state + + temp_adam_ui_state = UIState(self, adam_config) + window = MuonAdamWindow(self, self.train_config, temp_adam_ui_state, current_optimizer) + self.wait_window(window) + + self.train_config.optimizer.muon_adam_config = adam_config.to_dict() diff --git a/modules/ui/SampleWindowController.py b/modules/ui/SampleWindowController.py new file mode 100644 index 000000000..0f91ad2fa --- /dev/null +++ b/modules/ui/SampleWindowController.py @@ -0,0 +1,227 @@ +import contextlib +import copy +import os +import tkinter as tk +import traceback + +from modules.model.BaseModel import BaseModel +from modules.modelSampler.BaseModelSampler import ( + BaseModelSampler, + ModelSamplerOutput, +) +from modules.ui.SampleFrame import SampleFrame +from modules.util import create +from modules.util.callbacks.TrainCallbacks import TrainCallbacks +from modules.util.commands.TrainCommands import TrainCommands +from modules.util.config.SampleConfig import SampleConfig +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.EMAMode import EMAMode +from modules.util.enum.FileType import FileType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.time_util import get_string_timestamp +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import torch + +import customtkinter as ctk +from PIL import Image + + +class SampleWindow(ctk.CTkToplevel): + def __init__( + self, + parent, + train_config: TrainConfig, + use_external_model: bool, + callbacks: TrainCallbacks | None = None, + commands: TrainCommands | None = None, + *args, **kwargs + ): + super().__init__(parent, *args, **kwargs) + + self.title("Sample") + self.geometry("1200x800") + self.resizable(True, True) + + if not use_external_model: + self.initial_train_config = TrainConfig.default_values().from_dict(train_config.to_dict()) + # remove some settings to speed up model loading for sampling + self.initial_train_config.optimizer.optimizer = None + self.initial_train_config.ema = EMAMode.OFF + else: + self.initial_train_config = None + + #TODO why is there a current_train_config and an initial_train_config? + #current_train_config doesn't seem to ever change + self.current_train_config = train_config + self.callbacks = callbacks + self.commands = commands + + # get model specific defaults + model_type = train_config.model_type + self.sample = SampleConfig.default_values(model_type) + self.ui_state = UIState(self, self.sample) + + if use_external_model: + self.callbacks.set_on_sample_custom(self.__update_preview) + self.callbacks.set_on_update_sample_custom_progress(self.__update_progress) + else: + self.model = None + self.model_sampler = None + + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=1) + self.grid_rowconfigure(2, weight=0) + self.grid_rowconfigure(3, weight=0) + self.grid_columnconfigure(0, weight=0) + self.grid_columnconfigure(1, weight=1) + + prompt_frame = SampleFrame(self, self.sample, self.ui_state, include_settings=False, model_type=model_type) + prompt_frame.grid(row=0, column=0, columnspan=2, padx=0, pady=0, sticky="nsew") + + settings_frame = SampleFrame(self, self.sample, self.ui_state, include_prompt=False, model_type=model_type) + settings_frame.grid(row=1, column=0, padx=0, pady=0, sticky="nsew") + + # image + self.image = ctk.CTkImage( + light_image=self.__dummy_image(), + size=(512, 512) + ) + + image_label = ctk.CTkLabel(master=self, text="", image=self.image, height=512, width=512) + image_label.grid(row=1, column=1, rowspan=3, sticky="nsew") + + self.progress = components.progress(self, 2, 0) + components.button(self, 3, 0, "sample", self.__sample) + + self.wait_visibility() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + def __load_model(self) -> BaseModel: + model_loader = create.create_model_loader( + model_type=self.initial_train_config.model_type, + training_method=self.initial_train_config.training_method, + ) + + model_setup = create.create_model_setup( + model_type=self.initial_train_config.model_type, + train_device=torch.device(self.initial_train_config.train_device), + temp_device=torch.device(self.initial_train_config.temp_device), + training_method=self.initial_train_config.training_method, + ) + + model_names = self.initial_train_config.model_names() + if self.initial_train_config.continue_last_backup: + last_backup_path = self.initial_train_config.get_last_backup_path() + + if last_backup_path: + if self.initial_train_config.training_method == TrainingMethod.LORA: + model_names.lora = last_backup_path + elif self.initial_train_config.training_method == TrainingMethod.EMBEDDING: + model_names.embedding.model_name = last_backup_path + else: # fine-tunes + model_names.base_model = last_backup_path + + print(f"Loading from backup '{last_backup_path}'...") + else: + print("No backup found, loading without backup...") + + if self.initial_train_config.quantization.cache_dir is None: + self.initial_train_config.quantization.cache_dir = self.initial_train_config.cache_dir + "/quantization" + os.makedirs(self.initial_train_config.quantization.cache_dir, exist_ok=True) + + model = model_loader.load( + model_type=self.initial_train_config.model_type, + model_names=model_names, + weight_dtypes=self.initial_train_config.weight_dtypes(), + quantization=self.initial_train_config.quantization, + ) + model.train_config = self.initial_train_config + + model_setup.setup_optimizations(model, self.initial_train_config) + model_setup.setup_train_device(model, self.initial_train_config) + model_setup.setup_model(model, self.initial_train_config) + model.to(torch.device(self.initial_train_config.temp_device)) + + return model + + def __create_sampler(self, model: BaseModel) -> BaseModelSampler: + return create.create_model_sampler( + train_device=torch.device(self.initial_train_config.train_device), + temp_device=torch.device(self.initial_train_config.temp_device), + model=model, + model_type=self.initial_train_config.model_type, + training_method=self.initial_train_config.training_method, + ) + + def __update_preview(self, sampler_output: ModelSamplerOutput): + if sampler_output.file_type == FileType.IMAGE: + image = sampler_output.data + self.image.configure( + light_image=image, + size=(image.width, image.height), + ) + + def __update_progress(self, progress: int, max_progress: int): + self.progress.set(progress / max_progress) + self.update() + + def __dummy_image(self) -> Image: + return Image.new(mode="RGB", size=(512, 512), color=(0, 0, 0)) + + def __sample(self): + sample = copy.copy(self.sample) + + if self.commands: + self.commands.sample_custom(sample) + else: + if self.model is None: + # lazy initialization + self.model = self.__load_model() + self.model_sampler = self.__create_sampler(self.model) + + sample.from_train_config(self.current_train_config) + + sample_dir = os.path.join( + self.initial_train_config.workspace_dir, + "samples", + "custom", + ) + + progress = self.model.train_progress + sample_path = os.path.join( + sample_dir, + f"{get_string_timestamp()}-training-sample-{progress.filename_string()}" + ) + + self.model.eval() + + self.model_sampler.sample( + sample_config=sample, + destination=sample_path, + image_format=self.current_train_config.sample_image_format, + video_format=self.current_train_config.sample_video_format, + audio_format=self.current_train_config.sample_audio_format, + on_sample=self.__update_preview, + on_update_progress=self.__update_progress, + ) + + def destroy(self): + try: + if hasattr(self, "_icon_image_ref"): + del self._icon_image_ref + + # Remove any pending after callbacks + for after_id in self.tk.call('after', 'info'): + with contextlib.suppress(tk.TclError, RuntimeError): + self.after_cancel(after_id) + + super().destroy() + except (tk.TclError, RuntimeError) as e: + print(f"Error destroying window: {e}") + except Exception as e: + print(f"Unexpected error destroying window: {e}") + traceback.print_exc() diff --git a/modules/ui/TimestepDistributionWindowController.py b/modules/ui/TimestepDistributionWindowController.py new file mode 100644 index 000000000..21e41ce3e --- /dev/null +++ b/modules/ui/TimestepDistributionWindowController.py @@ -0,0 +1,186 @@ + +from modules.modelSetup.mixin.ModelSetupNoiseMixin import ( + ModelSetupNoiseMixin, +) +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.TimestepDistribution import TimestepDistribution +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import torch +from torch import Tensor + +import customtkinter as ctk +from customtkinter import AppearanceModeTracker, ThemeManager +from matplotlib import pyplot as plt +from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg + + +class TimestepGenerator(ModelSetupNoiseMixin): + + def __init__( + self, + timestep_distribution: TimestepDistribution, + min_noising_strength: float, + max_noising_strength: float, + noising_weight: float, + noising_bias: float, + timestep_shift: float, + ): + super().__init__() + + self.timestep_distribution = timestep_distribution + self.min_noising_strength = min_noising_strength + self.max_noising_strength = max_noising_strength + self.noising_weight = noising_weight + self.noising_bias = noising_bias + self.timestep_shift = timestep_shift + + def generate(self) -> Tensor: + generator = torch.Generator() + generator.seed() + + config = TrainConfig.default_values() + config.timestep_distribution = self.timestep_distribution + config.min_noising_strength = self.min_noising_strength + config.max_noising_strength = self.max_noising_strength + config.noising_weight = self.noising_weight + config.noising_bias = self.noising_bias + config.timestep_shift = self.timestep_shift + + + return self._get_timestep_discrete( + num_train_timesteps=1000, + deterministic=False, + generator=generator, + batch_size=1000000, + config=config, + ) + + +class TimestepDistributionWindow(ctk.CTkToplevel): + def __init__( + self, + parent, + config: TrainConfig, + ui_state: UIState, + *args, **kwargs, + ): + super().__init__(parent, *args, **kwargs) + + self.title("Timestep Distribution") + self.geometry("900x600") + self.resizable(True, True) + + self.config = config + self.ui_state = ui_state + self.image_preview_file_index = 0 + self.ax = None + self.canvas = None + + self.grid_rowconfigure(0, weight=1) + self.grid_columnconfigure(0, weight=1) + + frame = self.__content_frame(self) + frame.grid(row=0, column=0, sticky='nsew') + components.button(self, 1, 0, "ok", self.__ok) + + self.wait_visibility() + self.after(200, lambda: set_window_icon(self)) + self.grab_set() + self.focus_set() + + def __content_frame(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=0) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + frame.grid_rowconfigure(7, weight=1) + + # timestep distribution + components.label(frame, 0, 0, "Timestep Distribution", + tooltip="Selects the function to sample timesteps during training", + wide_tooltip=True) + components.options(frame, 0, 1, [str(x) for x in list(TimestepDistribution)], self.ui_state, + "timestep_distribution") + + # min noising strength + components.label(frame, 1, 0, "Min Noising Strength", + tooltip="Specifies the minimum noising strength used during training. This can help to improve composition, but prevents finer details from being trained") + components.entry(frame, 1, 1, self.ui_state, "min_noising_strength") + + # max noising strength + components.label(frame, 2, 0, "Max Noising Strength", + tooltip="Specifies the maximum noising strength used during training. This can be useful to reduce overfitting, but also reduces the impact of training samples on the overall image composition") + components.entry(frame, 2, 1, self.ui_state, "max_noising_strength") + + # noising weight + components.label(frame, 3, 0, "Noising Weight", + tooltip="Controls the weight parameter of the timestep distribution function. Use the preview to see more details.") + components.entry(frame, 3, 1, self.ui_state, "noising_weight") + + # noising bias + components.label(frame, 4, 0, "Noising Bias", + tooltip="Controls the bias parameter of the timestep distribution function. Use the preview to see more details.") + components.entry(frame, 4, 1, self.ui_state, "noising_bias") + + # timestep shift + components.label(frame, 5, 0, "Timestep Shift", + tooltip="Shift the timestep distribution. Use the preview to see more details.") + components.entry(frame, 5, 1, self.ui_state, "timestep_shift") + + # dynamic timestep shifting + components.label(frame, 6, 0, "Dynamic Timestep Shifting", + tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Dynamic Timestep Shifting is not shown in the preview. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) + components.switch(frame, 6, 1, self.ui_state, "dynamic_timestep_shifting") + + + # plot + appearance_mode = AppearanceModeTracker.get_mode() + background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) + text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) + background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" + text_color = f"#{int(text_color[0]/256):x}{int(text_color[1]/256):x}{int(text_color[2]/256):x}" + + fig, ax = plt.subplots() + self.ax = ax + self.canvas = FigureCanvasTkAgg(fig, master=frame) + self.canvas.get_tk_widget().grid(row=0, column=3, rowspan=8) + + fig.set_facecolor(background_color) + ax.set_facecolor(background_color) + ax.spines['bottom'].set_color(text_color) + ax.spines['left'].set_color(text_color) + ax.spines['top'].set_color(text_color) + ax.spines['right'].set_color(text_color) + ax.tick_params(axis='x', colors=text_color, which="both") + ax.tick_params(axis='y', colors=text_color, which="both") + ax.xaxis.label.set_color(text_color) + ax.yaxis.label.set_color(text_color) + + self.__update_preview() + + # update button + components.button(frame, 8, 3, "Update Preview", command=self.__update_preview) + + frame.pack(fill="both", expand=1) + return frame + + def __update_preview(self): + generator = TimestepGenerator( + timestep_distribution=self.config.timestep_distribution, + min_noising_strength=self.config.min_noising_strength, + max_noising_strength=self.config.max_noising_strength, + noising_weight=self.config.noising_weight, + noising_bias=self.config.noising_bias, + timestep_shift=self.config.timestep_shift, + ) + + self.ax.cla() + self.ax.hist(generator.generate(), bins=1000, range=(0, 999)) + self.canvas.draw() + + def __ok(self): + self.destroy() diff --git a/modules/ui/TopBarController.py b/modules/ui/TopBarController.py new file mode 100644 index 000000000..820fdb71a --- /dev/null +++ b/modules/ui/TopBarController.py @@ -0,0 +1,260 @@ +import json +import os +import traceback +import webbrowser +from collections.abc import Callable +from contextlib import suppress + +from modules.util import path_util +from modules.util.config.SecretsConfig import SecretsConfig +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ModelType import ModelType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.optimizer_util import change_optimizer +from modules.util.path_util import write_json_atomic +from modules.util.ui import components, dialogs +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class TopBar: + def __init__( + self, + master, + train_config: TrainConfig, + ui_state: UIState, + change_model_type_callback: Callable[[ModelType], None], + change_training_method_callback: Callable[[TrainingMethod], None], + load_preset_callback: Callable[[], None], + ): + self.master = master + self.train_config = train_config + self.ui_state = ui_state + self.change_model_type_callback = change_model_type_callback + self.change_training_method_callback = change_training_method_callback + self.load_preset_callback = load_preset_callback + + self.dir = "training_presets" + + self.config_ui_data = { + "config_name": path_util.canonical_join(self.dir, "#.json") + } + 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.current_config = [] + + self.frame = ctk.CTkFrame(master=master, corner_radius=0) + self.frame.grid(row=0, column=0, sticky="nsew") + + self.training_method = None + + # title + components.app_title(self.frame, 0, 0) + + # dropdown + self.configs_dropdown = None + self.__create_configs_dropdown() + + # remove button + # TODO + # components.icon_button(self.frame, 0, 2, "-", self.__remove_config) + + # Wiki button + components.button(self.frame, 0, 4, "Wiki", self.open_wiki, width=50) + + # save button + components.button(self.frame, 0, 3, "Save config", self.__save_config, + tooltip="Save the current configuration in a custom preset", width=90) + + # padding + self.frame.grid_columnconfigure(5, weight=1) + + # model type + components.options_kv( + master=self.frame, + row=0, + column=6, + values=[ #TODO simplify + ("SD1.5", ModelType.STABLE_DIFFUSION_15), + ("SD1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), + ("SD2.0", ModelType.STABLE_DIFFUSION_20), + ("SD2.0 Inpainting", ModelType.STABLE_DIFFUSION_20_INPAINTING), + ("SD2.1", ModelType.STABLE_DIFFUSION_21), + ("SD3", ModelType.STABLE_DIFFUSION_3), + ("SD3.5", ModelType.STABLE_DIFFUSION_35), + ("SDXL", ModelType.STABLE_DIFFUSION_XL_10_BASE), + ("SDXL Inpainting", ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING), + ("Wuerstchen v2", ModelType.WUERSTCHEN_2), + ("Stable Cascade", ModelType.STABLE_CASCADE_1), + ("PixArt Alpha", ModelType.PIXART_ALPHA), + ("PixArt Sigma", ModelType.PIXART_SIGMA), + ("Flux Dev.1", ModelType.FLUX_DEV_1), + ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), + ("Flux 2 [Dev, Klein]", ModelType.FLUX_2), + ("Sana", ModelType.SANA), + ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), + ("HiDream Full", ModelType.HI_DREAM_FULL), + ("Chroma1", ModelType.CHROMA_1), + ("QwenImage", ModelType.QWEN), + ("Z-Image", ModelType.Z_IMAGE), + ("Ernie Image", ModelType.ERNIE), + ], + ui_state=self.ui_state, + var_name="model_type", + command=self.__change_model_type, + ) + + def __create_training_method(self): + if self.training_method: + self.training_method.destroy() + + values = [] + #TODO simplify + if self.train_config.model_type.is_stable_diffusion(): + values = [ + ("Fine Tune", TrainingMethod.FINE_TUNE), + ("LoRA", TrainingMethod.LORA), + ("Embedding", TrainingMethod.EMBEDDING), + ("Fine Tune VAE", TrainingMethod.FINE_TUNE_VAE), + ] + elif self.train_config.model_type.is_stable_diffusion_3() \ + or self.train_config.model_type.is_stable_diffusion_xl() \ + or self.train_config.model_type.is_wuerstchen() \ + or self.train_config.model_type.is_pixart() \ + or self.train_config.model_type.is_flux_1() \ + or self.train_config.model_type.is_sana() \ + or self.train_config.model_type.is_hunyuan_video() \ + or self.train_config.model_type.is_hi_dream() \ + or self.train_config.model_type.is_chroma(): + values = [ + ("Fine Tune", TrainingMethod.FINE_TUNE), + ("LoRA", TrainingMethod.LORA), + ("Embedding", TrainingMethod.EMBEDDING), + ] + elif self.train_config.model_type.is_qwen() \ + or self.train_config.model_type.is_z_image() \ + or self.train_config.model_type.is_flux_2() \ + or self.train_config.model_type.is_ernie(): + values = [ + ("Fine Tune", TrainingMethod.FINE_TUNE), + ("LoRA", TrainingMethod.LORA), + ] + + # training method + self.training_method = components.options_kv( + master=self.frame, + row=0, + column=7, + values=values, + ui_state=self.ui_state, + var_name="training_method", + command=self.change_training_method_callback, + ) + + def __change_model_type(self, model_type: ModelType): + self.change_model_type_callback(model_type) + self.__create_training_method() + + def __create_configs_dropdown(self): + if self.configs_dropdown is not None: + self.configs_dropdown.grid_forget() + + self.configs_dropdown = components.options_kv( + self.frame, 0, 1, self.configs, self.config_ui_state, "config_name", self.__load_current_config + ) + + def __load_available_config_names(self): + if os.path.isdir(self.dir): + for path in os.listdir(self.dir): + if path != "#.json": + path = path_util.canonical_join(self.dir, path) + if path.endswith(".json") and os.path.isfile(path): + name = os.path.basename(path) + name = os.path.splitext(name)[0] + self.configs.append((name, path)) + self.configs.sort() + + def __save_to_file(self, name) -> str: + name = path_util.safe_filename(name) + path = path_util.canonical_join("training_presets", f"{name}.json") + + write_json_atomic(path, self.train_config.to_settings_dict(secrets=False)) + + return path + + def __save_secrets(self, path) -> str: + write_json_atomic(path, self.train_config.secrets.to_dict()) + return path + + def open_wiki(self): + webbrowser.open("https://github.com/Nerogar/OneTrainer/wiki", new=0, autoraise=False) + + def __save_new_config(self, name): + path = self.__save_to_file(name) + + is_new_config = name not in [x[0] for x in self.configs] + + if is_new_config: + self.configs.append((name, path)) + self.configs.sort() + + if self.config_ui_data["config_name"] != path_util.canonical_join(self.dir, f"{name}.json"): + self.config_ui_state.get_var("config_name").set(path_util.canonical_join(self.dir, f"{name}.json")) + + if is_new_config: + self.__create_configs_dropdown() + + def __save_config(self): + default_value = self.configs_dropdown.get() + while default_value.startswith('#'): + default_value = default_value[1:] + + dialogs.StringInputDialog( + parent=self.master, + title="name", + question="Config Name", + callback=self.__save_new_config, + default_value=default_value, + validate_callback=lambda x: not x.startswith("#") + ) + + def __load_current_config(self, filename): + try: + basename = os.path.basename(filename) + is_built_in_preset = basename.startswith("#") and basename != "#.json" + + with open(filename, "r") as f: + loaded_dict = json.load(f) + default_config = TrainConfig.default_values() + if is_built_in_preset: + # always assume built-in configs are saved in the most recent version + loaded_dict["__version"] = default_config.config_version + loaded_config = default_config.from_dict(loaded_dict).to_unpacked_config() + + with suppress(FileNotFoundError), open("secrets.json", "r") as f: + secrets_dict=json.load(f) + loaded_config.secrets = SecretsConfig.default_values().from_dict(secrets_dict) + + self.train_config.from_dict(loaded_config.to_dict()) + self.ui_state.update(loaded_config) + + optimizer_config = change_optimizer(self.train_config) + self.ui_state.get_var("optimizer").update(optimizer_config) + + self.load_preset_callback() + except FileNotFoundError: + pass + except Exception: + print(traceback.format_exc()) + + def __remove_config(self): + # TODO + pass + + def save_default(self): + self.__save_to_file("#") + self.__save_secrets("secrets.json") diff --git a/modules/ui/TrainUIController.py b/modules/ui/TrainUIController.py new file mode 100644 index 000000000..ba90d2e64 --- /dev/null +++ b/modules/ui/TrainUIController.py @@ -0,0 +1,889 @@ +import ctypes +import datetime +import json +import os +import platform +import subprocess +import sys +import threading +import time +import traceback +import webbrowser +from collections.abc import Callable +from contextlib import suppress +from pathlib import Path +from tkinter import filedialog, messagebox + +import scripts.generate_debug_report +from modules.ui.AdditionalEmbeddingsTab import AdditionalEmbeddingsTab +from modules.ui.CaptionUI import CaptionUI +from modules.ui.CloudTab import CloudTab +from modules.ui.ConceptTab import ConceptTab +from modules.ui.ConvertModelUI import ConvertModelUI +from modules.ui.LoraTab import LoraTab +from modules.ui.ModelTab import ModelTab +from modules.ui.ProfilingWindow import ProfilingWindow +from modules.ui.SampleWindow import SampleWindow +from modules.ui.SamplingTab import SamplingTab +from modules.ui.TopBar import TopBar +from modules.ui.TrainingTab import TrainingTab +from modules.ui.VideoToolUI import VideoToolUI +from modules.util import create +from modules.util.callbacks.TrainCallbacks import TrainCallbacks +from modules.util.commands.TrainCommands import TrainCommands +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.DataType import DataType +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.PathIOType import PathIOType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.torch_util import torch_gc +from modules.util.TrainProgress import TrainProgress +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState +from modules.util.ui.validation import flush_and_validate_all + +import torch + +import customtkinter as ctk +from customtkinter import AppearanceModeTracker + +# 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": + with suppress(Exception): + # 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): + set_step_progress: Callable[[int, int], None] + set_epoch_progress: Callable[[int, int], None] + + status_label: ctk.CTkLabel | None + training_button: ctk.CTkButton | None + training_callbacks: TrainCallbacks | None + training_commands: TrainCommands | None + + _TRAIN_BUTTON_STYLES = { + "idle": { + "text": "Start Training", + "state": "normal", + "fg_color": "#198754", + "hover_color": "#146c43", + "text_color": "white", + "text_color_disabled": "white", + }, + "running": { + "text": "Stop Training", + "state": "normal", + "fg_color": "#dc3545", + "hover_color": "#bb2d3b", + "text_color": "white", + }, + "stopping": { + "text": "Stopping...", + "state": "disabled", + "fg_color": "#dc3545", + "hover_color": "#dc3545", + "text_color": "white", + "text_color_disabled": "white", + }, + } + + def __init__(self): + super().__init__() + + self.title("OneTrainer") + self.geometry("1100x740") + + self.after(100, lambda: self._set_icon()) + + # more efficient version of ctk.set_appearance_mode("System"), which retrieves the system theme on each main loop iteration + ctk.set_appearance_mode("Light" if AppearanceModeTracker.detect_appearance_mode() == 0 else "Dark") + ctk.set_default_color_theme("blue") + + self.train_config = TrainConfig.default_values() + self.ui_state = UIState(self, self.train_config) + + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=1) + self.grid_rowconfigure(2, weight=0) + self.grid_columnconfigure(0, weight=1) + + self.status_label = None + self.eta_label = None + self.training_button = None + self.export_button = None + self.tabview = None + + self.model_tab = None + self.training_tab = None + self.lora_tab = None + self.cloud_tab = None + self.additional_embeddings_tab = None + + self.top_bar_component = self.top_bar(self) + self.content_frame(self) + self.bottom_bar(self) + + 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.workspace_dir_trace_id = self.ui_state.add_var_trace("workspace_dir", self._on_workspace_dir_change_trace) + + # Persistent profiling window. + self.profiling_window = ProfilingWindow(self) + + self.protocol("WM_DELETE_WINDOW", self.__close) + + def __close(self): + self.top_bar_component.save_default() + self._stop_always_on_tensorboard() + if hasattr(self, 'workspace_dir_trace_id'): + self.ui_state.remove_var_trace("workspace_dir", self.workspace_dir_trace_id) + self.quit() + + def top_bar(self, master): + return TopBar( + master, + self.train_config, + self.ui_state, + self.change_model_type, + self.change_training_method, + self.load_preset, + ) + + def _set_icon(self): + """Set the window icon safely after window is ready""" + set_window_icon(self) + + def bottom_bar(self, master): + frame = ctk.CTkFrame(master=master, corner_radius=0) + frame.grid(row=2, column=0, sticky="nsew") + + self.set_step_progress, self.set_epoch_progress = components.double_progress(frame, 0, 0, "step", "epoch") + + # status + ETA container + self.status_frame = ctk.CTkFrame(frame, corner_radius=0, fg_color="transparent") + self.status_frame.grid(row=0, column=1, sticky="w") + self.status_frame.grid_rowconfigure(0, weight=0) + self.status_frame.grid_rowconfigure(1, weight=0) + self.status_frame.grid_columnconfigure(0, weight=1) + + self.status_label = components.label(self.status_frame, 0, 0, "", pad=0, + tooltip="Current status of the training run") + self.eta_label = components.label(self.status_frame, 1, 0, "", pad=0) + + # padding + frame.grid_columnconfigure(2, weight=1) + + + # export button + self.export_button = components.button(frame, 0, 3, "Export", self.export_training, + width=60, padx=5, pady=(15, 0), + tooltip="Export the current configuration as a script to run without a UI") + + # debug button + components.button(frame, 0, 4, "Debug", self.generate_debug_package, + width=60, padx=(5, 25), pady=(15, 0), + tooltip="Generate a zip file with config.json, debug_report.log and settings diff, use this to report bugs or issues") + + # tensorboard button + components.button(frame, 0, 5, "Tensorboard", self.open_tensorboard, + width=100, padx=(0, 5), pady=(15, 0)) + + # training button + self.training_button = components.button(frame, 0, 6, "Start Training", self.start_training, + padx=(5, 20), pady=(15, 0)) + self._set_training_button_style("idle") # centralized styling + + return frame + + def content_frame(self, master): + frame = ctk.CTkFrame(master=master, corner_radius=0) + frame.grid(row=1, column=0, sticky="nsew") + + frame.grid_rowconfigure(0, weight=1) + frame.grid_columnconfigure(0, weight=1) + + self.tabview = ctk.CTkTabview(frame) + self.tabview.grid(row=0, column=0, sticky="nsew") + + 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) + + return frame + + def create_general_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # workspace dir + components.label(frame, 0, 0, "Workspace Directory", + tooltip="The directory where all files of this training run are saved") + components.path_entry(frame, 0, 1, self.ui_state, "workspace_dir", mode="dir", command=self._on_workspace_dir_change) + + # cache dir + components.label(frame, 0, 2, "Cache Directory", + tooltip="The directory where cached data is saved") + components.path_entry(frame, 0, 3, self.ui_state, "cache_dir", mode="dir") + + # continue from previous backup + components.label(frame, 2, 0, "Continue from last backup", + tooltip="Automatically continues training from the last backup saved in /backup") + components.switch(frame, 2, 1, self.ui_state, "continue_last_backup") + + # only cache + components.label(frame, 2, 2, "Only Cache", + tooltip="Only populate the cache, without any training") + components.switch(frame, 2, 3, self.ui_state, "only_cache") + + # TODO: In Phase 4 rework the general tab. + # prevent overwrites + components.label(frame, 3, 0, "Prevent Overwrites", + tooltip="When enabled, output paths that already exist on disk will be flagged as invalid to avoid accidental overwrites") + components.switch(frame, 3, 1, self.ui_state, "prevent_overwrites") + + # debug + components.label(frame, 4, 0, "Debug mode", + tooltip="Save debug information during the training into the debug directory") + components.switch(frame, 4, 1, self.ui_state, "debug_mode") + + components.label(frame, 4, 2, "Debug Directory", + tooltip="The directory where debug data is saved") + components.path_entry(frame, 4, 3, self.ui_state, "debug_dir", mode="dir", io_type=PathIOType.OUTPUT) + + # 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", + tooltip="Port to use for Tensorboard link") + components.entry(frame, 7, 3, self.ui_state, "tensorboard_port") + + + # validation + components.label(frame, 8, 0, "Validation", + tooltip="Enable validation steps and add new graph in tensorboard") + components.switch(frame, 8, 1, self.ui_state, "validation") + + components.label(frame, 8, 2, "Validate after", + tooltip="The interval used when validate training") + components.time_entry(frame, 8, 3, self.ui_state, "validate_after", "validate_after_unit") + + # device + components.label(frame, 10, 0, "Dataloader Threads", + tooltip="Number of threads used for the data loader. Increase if your GPU has room during caching, decrease if it's going out of memory during caching.") + components.entry(frame, 10, 1, self.ui_state, "dataloader_threads", required=True) + + components.label(frame, 11, 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", required=True) + + components.label(frame, 12, 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", + 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.label(frame, 13, 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, + "gradient_reduce_precision") + + components.label(frame, 13, 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.label(frame, 14, 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)", + 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.label(frame, 15, 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") + + frame.pack(fill="both", expand=1) + return frame + + def create_model_tab(self, master): + return ModelTab(master, self.train_config, self.ui_state) + + def create_data_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + + # aspect ratio bucketing + components.label(frame, 0, 0, "Aspect Ratio Bucketing", + tooltip="Aspect ratio bucketing enables training on images with different aspect ratios") + components.switch(frame, 0, 1, self.ui_state, "aspect_ratio_bucketing") + + # latent caching + components.label(frame, 1, 0, "Latent Caching", + tooltip="Caching of intermediate training data that can be re-used between epochs") + components.switch(frame, 1, 1, self.ui_state, "latent_caching") + + # clear cache before training + components.label(frame, 2, 0, "Clear cache before training", + tooltip="Clears the cache directory before starting to train. Only disable this if you want to continue using the same cached data. Disabling this can lead to errors, if other settings are changed during a restart") + components.switch(frame, 2, 1, self.ui_state, "clear_cache_before_training") + + frame.pack(fill="both", expand=1) + return frame + + def create_concepts_tab(self, master): + return ConceptTab(master, self.train_config, self.ui_state) + + def create_training_tab(self, master) -> TrainingTab: + return TrainingTab(master, self.train_config, self.ui_state) + + def create_cloud_tab(self, master) -> CloudTab: + return CloudTab(master, self.train_config, self.ui_state,parent=self) + + def create_sampling_tab(self, master): + master.grid_rowconfigure(0, weight=0) + master.grid_rowconfigure(1, weight=1) + master.grid_columnconfigure(0, weight=1) + + # sample after + top_frame = ctk.CTkFrame(master=master, corner_radius=0) + top_frame.grid(row=0, column=0, sticky="nsew") + sub_frame = ctk.CTkFrame(master=top_frame, corner_radius=0, fg_color="transparent") + sub_frame.grid(row=1, column=0, sticky="nsew", columnspan=6) + + components.label(top_frame, 0, 0, "Sample After", + tooltip="The interval used when automatically sampling from the model during training") + components.time_entry(top_frame, 0, 1, self.ui_state, "sample_after", "sample_after_unit") + + components.label(top_frame, 0, 2, "Skip First", + tooltip="Start sampling automatically after this interval has elapsed.") + components.entry(top_frame, 0, 3, self.ui_state, "sample_skip_first", width=50, sticky="nw") + + components.label(top_frame, 0, 4, "Format", + tooltip="File Format used when saving samples") + components.options_kv(top_frame, 0, 5, [ + ("PNG", ImageFormat.PNG), + ("JPG", ImageFormat.JPG), + ], self.ui_state, "sample_image_format") + + components.button(top_frame, 0, 6, "sample now", self.sample_now) + + components.button(top_frame, 0, 7, "manual sample", self.open_manual_sample_window ) + + components.label(sub_frame, 0, 0, "Non-EMA Sampling", + tooltip="Whether to include non-ema sampling when using ema.") + components.switch(sub_frame, 0, 1, self.ui_state, "non_ema_sampling") + + components.label(sub_frame, 0, 2, "Samples to Tensorboard", + tooltip="Whether to include sample images in the Tensorboard output.") + components.switch(sub_frame, 0, 3, self.ui_state, "samples_to_tensorboard") + + # table + frame = ctk.CTkFrame(master=master, corner_radius=0) + frame.grid(row=1, column=0, sticky="nsew") + + return SamplingTab(frame, self.train_config, self.ui_state) + + def create_backup_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + + # backup after + components.label(frame, 0, 0, "Backup After", + tooltip="The interval used when automatically creating model backups during training") + components.time_entry(frame, 0, 1, self.ui_state, "backup_after", "backup_after_unit") + + # backup now + components.button(frame, 0, 3, "backup now", self.backup_now) + + # rolling backup + components.label(frame, 1, 0, "Rolling Backup", + tooltip="If rolling backups are enabled, older backups are deleted automatically") + components.switch(frame, 1, 1, self.ui_state, "rolling_backup") + + # rolling backup count + components.label(frame, 1, 3, "Rolling Backup Count", + tooltip="Defines the number of backups to keep if rolling backups are enabled") + components.entry(frame, 1, 4, self.ui_state, "rolling_backup_count") + + # backup before save + components.label(frame, 2, 0, "Backup Before Save", + tooltip="Create a full backup before saving the final model") + components.switch(frame, 2, 1, self.ui_state, "backup_before_save") + + # save after + components.label(frame, 3, 0, "Save Every", + tooltip="The interval used when automatically saving the model during training") + components.time_entry(frame, 3, 1, self.ui_state, "save_every", "save_every_unit") + + # save now + components.button(frame, 3, 3, "save now", self.save_now) + + # skip save + components.label(frame, 4, 0, "Skip First", + tooltip="Start saving automatically after this interval has elapsed") + components.entry(frame, 4, 1, self.ui_state, "save_skip_first", width=50, sticky="nw") + + # save filename prefix + components.label(frame, 5, 0, "Save Filename Prefix", + tooltip="The prefix for filenames used when saving the model during training") + components.entry(frame, 5, 1, self.ui_state, "save_filename_prefix") + + frame.pack(fill="both", expand=1) + return frame + + def embedding_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + + # 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.path_entry( + frame, 0, 1, self.ui_state, "embedding.model_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # token count + components.label(frame, 1, 0, "Token count", + tooltip="The token count used when creating a new embedding. Leave empty to auto detect from the initial embedding text.") + components.entry(frame, 1, 1, self.ui_state, "embedding.token_count") + + # initial embedding text + components.label(frame, 2, 0, "Initial embedding text", + tooltip="The initial embedding text used when creating a new embedding") + components.entry(frame, 2, 1, self.ui_state, "embedding.initial_embedding_text") + + # embedding weight dtype + components.label(frame, 3, 0, "Embedding Weight Data Type", + tooltip="The Embedding weight data type used for training. This can reduce memory consumption, but reduces precision") + components.options_kv(frame, 3, 1, [ + ("float32", DataType.FLOAT_32), + ("bfloat16", DataType.BFLOAT_16), + ], self.ui_state, "embedding_weight_dtype") + + # placeholder + components.label(frame, 4, 0, "Placeholder", + tooltip="The placeholder used when using the embedding in a prompt") + components.entry(frame, 4, 1, self.ui_state, "embedding.placeholder") + + # output embedding + components.label(frame, 5, 0, "Output embedding", + tooltip="Output embeddings are calculated at the output of the text encoder, not the input. This can improve results for larger text encoders and lower VRAM usage.") + components.switch(frame, 5, 1, self.ui_state, "embedding.is_output_embedding") + + frame.pack(fill="both", expand=1) + return frame + + def create_additional_embeddings_tab(self, master): + return AdditionalEmbeddingsTab(master, self.train_config, self.ui_state) + + def create_tools_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + + # dataset + components.label(frame, 0, 0, "Dataset Tools", + tooltip="Open the captioning tool") + components.button(frame, 0, 1, "Open", self.open_dataset_tool) + + # video tools + components.label(frame, 1, 0, "Video Tools", + tooltip="Open the video tools") + components.button(frame, 1, 1, "Open", self.open_video_tool) + + # convert model + components.label(frame, 2, 0, "Convert Model Tools", + tooltip="Open the model conversion tool") + components.button(frame, 2, 1, "Open", self.open_convert_model_tool) + + # sample + components.label(frame, 3, 0, "Sampling Tool", + tooltip="Open the model sampling tool") + components.button(frame, 3, 1, "Open", self.open_sampling_tool) + + components.label(frame, 4, 0, "Profiling Tool", + tooltip="Open the profiling tools.") + components.button(frame, 4, 1, "Open", self.open_profiling_tool) + + frame.pack(fill="both", expand=1) + return frame + + def change_model_type(self, model_type: ModelType): + if self.model_tab: + self.model_tab.refresh_ui() + + if self.training_tab: + self.training_tab.refresh_ui() + + if self.lora_tab: + self.lora_tab.refresh_ui() + + def change_training_method(self, training_method: TrainingMethod): + if not self.tabview: + return + + if self.model_tab: + self.model_tab.refresh_ui() + + if training_method != TrainingMethod.LORA and "LoRA" in self.tabview._tab_dict: + self.tabview.delete("LoRA") + self.lora_tab = None + if training_method != TrainingMethod.EMBEDDING and "embedding" in self.tabview._tab_dict: + self.tabview.delete("embedding") + + if training_method == TrainingMethod.LORA and "LoRA" not in self.tabview._tab_dict: + self.lora_tab = LoraTab(self.tabview.add("LoRA"), self.train_config, self.ui_state) + if training_method == TrainingMethod.EMBEDDING and "embedding" not in self.tabview._tab_dict: + self.embedding_tab(self.tabview.add("embedding")) + + def load_preset(self): + if not self.tabview: + return + + if self.additional_embeddings_tab: + self.additional_embeddings_tab.refresh_ui() + + def open_tensorboard(self): + webbrowser.open("http://localhost:" + str(self.train_config.tensorboard_port), new=0, autoraise=False) + + def _calculate_eta_string(self, train_progress: TrainProgress, max_step: int, max_epoch: int) -> str | None: + spent_total = time.monotonic() - self.start_time + steps_done = train_progress.epoch * max_step + train_progress.epoch_step + remaining_steps = (max_epoch - train_progress.epoch - 1) * max_step + (max_step - train_progress.epoch_step) + total_eta = spent_total / steps_done * remaining_steps + + if train_progress.global_step <= 30: + return "Estimating ..." + + td = datetime.timedelta(seconds=total_eta) + days = td.days + hours, remainder = divmod(td.seconds, 3600) + minutes, seconds = divmod(remainder, 60) + if days > 0: + return f"{days}d {hours}h" + elif hours > 0: + return f"{hours}h {minutes}m" + elif minutes > 0: + return f"{minutes}m {seconds}s" + else: + return f"{seconds}s" + + def set_eta_label(self, train_progress: TrainProgress, max_step: int, max_epoch: int): + eta_str = self._calculate_eta_string(train_progress, max_step, max_epoch) + if eta_str is not None: + self.eta_label.configure(text=f"ETA: {eta_str}") + else: + self.eta_label.configure(text="") + + def delete_eta_label(self): + self.eta_label.configure(text="") + + def on_update_train_progress(self, train_progress: TrainProgress, max_step: int, max_epoch: int): + self.set_step_progress(train_progress.epoch_step, max_step) + self.set_epoch_progress(train_progress.epoch, max_epoch) + self.set_eta_label(train_progress, max_step, max_epoch) + + def on_update_status(self, status: str): + self.status_label.configure(text=status) + + def open_dataset_tool(self): + window = CaptionUI(self, None, False) + self.wait_window(window) + + def open_video_tool(self): + window = VideoToolUI(self) + self.wait_window(window) + + def open_convert_model_tool(self): + window = ConvertModelUI(self) + self.wait_window(window) + + def open_sampling_tool(self): + if not self.training_callbacks and not self.training_commands: + window = SampleWindow( + self, + use_external_model=False, + train_config=self.train_config, + ) + self.wait_window(window) + torch_gc() + + def open_profiling_tool(self): + self.profiling_window.deiconify() + + def generate_debug_package(self): + zip_path = filedialog.askdirectory( + initialdir=".", + title="Select Directory to Save Debug Package" + ) + + if not zip_path: + return + + zip_path = Path(zip_path) / "OneTrainer_debug_report.zip" + + self.on_update_status("Generating debug package...") + + try: + config_json_string = json.dumps(self.train_config.to_pack_dict(secrets=False)) + scripts.generate_debug_report.create_debug_package(str(zip_path), config_json_string) + self.on_update_status(f"Debug package saved to {zip_path.name}") + except Exception as e: + traceback.print_exc() + self.on_update_status(f"Error generating debug package: {e}") + + + def open_manual_sample_window (self): + training_callbacks = self.training_callbacks + training_commands = self.training_commands + + if training_callbacks and training_commands: + window = SampleWindow( + self, + train_config=self.train_config, + use_external_model=True, + callbacks=training_callbacks, + commands=training_commands, + ) + self.wait_window(window) + training_callbacks.set_on_sample_custom() + + def __training_thread_function(self): + error_caught = False + + self.training_callbacks = TrainCallbacks( + on_update_train_progress=self.on_update_train_progress, + on_update_status=self.on_update_status, + ) + + trainer = create.create_trainer(self.train_config, self.training_callbacks, self.training_commands, reattach=self.cloud_tab.reattach) + try: + trainer.start() + if self.train_config.cloud.enabled: + self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) + + self.start_time = time.monotonic() + trainer.train() + except Exception: + if self.train_config.cloud.enabled: + self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) + error_caught = True + traceback.print_exc() + + trainer.end() + + # clear gpu memory + del trainer + + 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() + + # 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) + + def start_training(self): + if self.training_thread is None: + self.save_default() + + # --- pre-training validation gate --- + errors = flush_and_validate_all() + + if errors: + bullet_list = "\n".join(f"• {e}" for e in errors) + messagebox.showerror( + "Cannot Start Training", + f"Please fix the following errors before training:\n\n{bullet_list}", + ) + return + + 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() + + self.training_thread = threading.Thread(target=self.__training_thread_function) + self.training_thread.start() + else: + self._set_training_button_stopping() + self.on_update_status("Stopping ...") + self.training_commands.stop() + + 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() + + def export_training(self): + file_path = filedialog.asksaveasfilename(filetypes=[ + ("All Files", "*.*"), + ("json", "*.json"), + ], initialdir=".", initialfile="config.json") + + if file_path: + with open(file_path, "w") as f: + json.dump(self.train_config.to_pack_dict(secrets=False), f, indent=4) + + def sample_now(self): + train_commands = self.training_commands + if train_commands: + train_commands.sample_default() + + def backup_now(self): + train_commands = self.training_commands + if train_commands: + train_commands.backup() + + def save_now(self): + train_commands = self.training_commands + 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") + + 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", + ] + + if self.train_config.tensorboard_expose: + tensorboard_args.append("--bind_all") + + 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 _set_training_button_style(self, mode: str): + if not self.training_button: + return + style = self._TRAIN_BUTTON_STYLES.get(mode) + if not style: + return + self.training_button.configure(**style) + + def _set_training_button_idle(self): + self._set_training_button_style("idle") + + def _set_training_button_running(self): + self._set_training_button_style("running") + + def _set_training_button_stopping(self): + self._set_training_button_style("stopping") diff --git a/modules/ui/VideoToolUIController.py b/modules/ui/VideoToolUIController.py new file mode 100644 index 000000000..c3291e6ea --- /dev/null +++ b/modules/ui/VideoToolUIController.py @@ -0,0 +1,877 @@ +import concurrent.futures +import math +import os +import pathlib +import random +import shlex +import subprocess +import threading +import webbrowser +from fractions import Fraction +from tkinter import filedialog + +from modules.util.image_util import load_image +from modules.util.path_util import SUPPORTED_VIDEO_EXTENSIONS +from modules.util.ui import components + +import av +import customtkinter as ctk +import cv2 +import scenedetect +from PIL import Image + + +class VideoToolUI(ctk.CTkToplevel): + def __init__( + self, + parent, + *args, **kwargs, + ): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + + self.title("Video Tools") + self.geometry("600x720") + self.resizable(True, True) + self.wait_visibility() + self.focus_set() + + self.grid_rowconfigure(0, weight=1) + self.grid_rowconfigure(1, weight=0) + self.grid_columnconfigure(0, weight=1) + + tabview = ctk.CTkTabview(self) + tabview.grid(row=0, column=0, sticky="nsew") + + self.clip_extract_tab = self.__clip_extract_tab(tabview.add("extract clips")) + self.image_extract_tab = self.__image_extract_tab(tabview.add("extract images")) + self.video_download_tab = self.__video_download_tab(tabview.add("download")) + self.status_bar(self) + + def status_bar(self, master): + frame = ctk.CTkFrame(master, fg_color="transparent") + frame.grid(row=1, column=0) + frame.grid_columnconfigure(0, weight=0, minsize=160) + frame.grid_columnconfigure(1, weight=0, minsize=300) + frame.grid_columnconfigure(2, weight=1) + + #create preview image + preview_path = "resources/icons/icon.png" + preview = load_image(preview_path, 'RGB') + preview.thumbnail((150, 150)) + self.preview_image= ctk.CTkImage(light_image=preview, size=preview.size) + self.preview_image_label = ctk.CTkLabel( + master=frame, text="Preview image", image=self.preview_image, height=150, width=150, + compound="top") + self.preview_image_label.grid(row=0, column=0, sticky="nw", padx=5, pady=5) + + #displays progress and messages that also go to terminal + self.status_label = ctk.CTkTextbox(master=frame, width=400, height=160, wrap="word", border_width=2) + self.status_label.insert(index="1.0", text="Current status") + self.status_label.configure(state="disabled") + self.status_label.grid(row=0, column=1, sticky="ne", padx=5, pady=5) + + def __clip_extract_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0, minsize=120) + frame.grid_columnconfigure(1, weight=0, minsize=200) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # single video + components.label(frame, 0, 0, "Single Video", + tooltip="Link to single video file to process.") + self.clip_single_entry = ctk.CTkEntry(frame, width=190) + self.clip_single_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) + self.clip_single_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_file(self.clip_single_entry, + [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))] + )) + self.clip_single_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) + components.button(frame, 0, 2, "Extract Single", + command=lambda: self.__extract_clips_button(False)) + + # time range + components.label(frame, 1, 0, " Time Range", + tooltip="Time range to limit selection for single video, \ + format as hour:minute:second, minute:second, or seconds.") + self.clip_time_start_entry = ctk.CTkEntry(frame, width=100) + self.clip_time_start_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.clip_time_start_entry.insert(0, "00:00:00") + self.clip_time_end_entry = ctk.CTkEntry(frame, width=100) + self.clip_time_end_entry.grid(row=1, column=1, sticky="e", padx=5, pady=5) + self.clip_time_end_entry.insert(0, "99:99:99") + + # directory of videos + components.label(frame, 2, 0, "Directory", + tooltip="Path to directory with multiple videos to process, including in subdirectories.") + self.clip_list_entry = ctk.CTkEntry(frame, width=190) + self.clip_list_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + self.clip_list_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_dir(self.clip_list_entry)) + self.clip_list_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) + components.button(frame, 2, 2, "Extract Directory", + command=lambda: self.__extract_clips_button(True)) + + # output directory + components.label(frame, 3, 0, "Output", + tooltip="Path to folder where extracted clips will be saved.") + self.clip_output_entry = ctk.CTkEntry(frame, width=190) + self.clip_output_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) + self.clip_output_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_dir(self.clip_output_entry)) + self.clip_output_button.grid(row=3, column=1, sticky="e", padx=5, pady=5) + + # output to subdirectories + self.output_subdir_clip = ctk.BooleanVar(self, False) + components.label(frame, 4, 0, "Output to\nSubdirectories", + tooltip="If enabled, files are saved to subfolders based on filename and input directory. \ + Otherwise will all be saved to the top level of the output directory.") + self.output_subdir_clip_entry = ctk.CTkSwitch(frame, variable=self.output_subdir_clip, text="") + self.output_subdir_clip_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + + # split at cuts + self.split_at_cuts = ctk.BooleanVar(self, False) + components.label(frame, 5, 0, "Split at Cuts", + tooltip="If enabled, detect cuts in the input video and split at those points. \ + Otherwise will split at any point, and clips may contain cuts.") + self.split_cuts_entry = ctk.CTkSwitch(frame, variable=self.split_at_cuts, text="") + self.split_cuts_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) + + # maximum length + components.label(frame, 6, 0, "Max Length (s)", + tooltip="Maximum length in seconds for saved clips, larger clips will be broken into multiple small clips.") + self.clip_length_entry = ctk.CTkEntry(frame, width=220) + self.clip_length_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) + self.clip_length_entry.insert(0, "3") + + # Set FPS + components.label(frame, 7, 0, "Set FPS", + tooltip="FPS to convert output videos to, set to 0 to keep original rate.") + self.clip_fps_entry = ctk.CTkEntry(frame, width=220) + self.clip_fps_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) + self.clip_fps_entry.insert(0, "24.0") + + # Remove borders + self.clip_bordercrop = ctk.BooleanVar(self, False) + components.label(frame, 8, 0, "Remove Borders", + tooltip="Remove black borders from output clip") + self.clip_bordercrop_entry = ctk.CTkSwitch(frame, variable=self.clip_bordercrop, text="") + self.clip_bordercrop_entry.grid(row=8, column=1, sticky="w", padx=5, pady=5) + + # Crop Variation + components.label(frame, 9, 0, "Crop Variation", + tooltip="Output clips will be randomly cropped to +- the base aspect ratio, \ + somewhat biased towards making square videos. Set to 0 to use only base aspect.") + self.clip_crop_entry = ctk.CTkEntry(frame, width=220) + self.clip_crop_entry.grid(row=9, column=1, sticky="w", padx=5, pady=5) + self.clip_crop_entry.insert(0, "0.2") + + # object filter - currently unused, may implement in future + # components.label(frame, 9, 0, "Object Filter", + # tooltip="Detect general features using Haar-Cascade classifier, and choose how to deal with clips where it is detected") + # components.options(frame, 9, 1, ["NONE", "FACE", "EYE", "BODY"], self.video_ui_state, "filter_object") + # components.options(frame, 9, 2, ["INCLUDE", "EXCLUDE", "SUBFOLDER"], self.video_ui_state, "filter_behavior") + + frame.pack(fill="both", expand=1) + return frame + + def __image_extract_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0, minsize=120) + frame.grid_columnconfigure(1, weight=0, minsize=200) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # single video + components.label(frame, 0, 0, "Single Video", + tooltip="Link to single video file to process.") + self.image_single_entry = ctk.CTkEntry(frame, width=190) + self.image_single_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) + self.image_single_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_file(self.image_single_entry, + [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))] + )) + self.image_single_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) + components.button(frame, 0, 2, "Extract Single", + command=lambda: self.__extract_images_button(False)) + + # time range + components.label(frame, 1, 0, " Time Range", + tooltip="Time range to limit selection for single video, \ + format as hour:minute:second, minute:second, or seconds.") + self.image_time_start_entry = ctk.CTkEntry(frame, width=100) + self.image_time_start_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.image_time_start_entry.insert(0, "00:00:00") + self.image_time_end_entry = ctk.CTkEntry(frame, width=100) + self.image_time_end_entry.grid(row=1, column=1, sticky="e", padx=5, pady=5) + self.image_time_end_entry.insert(0, "99:99:99") + + # directory of videos + components.label(frame, 2, 0, "Directory", + tooltip="Path to directory with multiple videos to process, including in subdirectories.") + self.image_list_entry = ctk.CTkEntry(frame, width=190) + self.image_list_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + self.image_list_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_dir(self.image_list_entry)) + self.image_list_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) + components.button(frame, 2, 2, "Extract Directory", + command=lambda: self.__extract_images_button(True)) + + # output directory + components.label(frame, 3, 0, "Output", + tooltip="Path to folder where extracted images will be saved.") + self.image_output_entry = ctk.CTkEntry(frame, width=190) + self.image_output_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) + self.image_output_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_dir(self.image_output_entry)) + self.image_output_button.grid(row=3, column=1, sticky="e", padx=5, pady=5) + + # output to subdirectories + self.output_subdir_img = ctk.BooleanVar(self, False) + components.label(frame, 4, 0, "Output to\nSubdirectories", + tooltip="If enabled, files are saved to subfolders based on filename and input directory. \ + Otherwise will all be saved to the top level of the output directory.") + self.output_subdir_img_entry = ctk.CTkSwitch(frame, variable=self.output_subdir_img, text="") + self.output_subdir_img_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + + # image capture rate + components.label(frame, 5, 0, "Images/sec", + tooltip="Number of images to capture per second of video. \ + Images will be taken at semi-random frames around the specified frequency.") + self.capture_rate_entry = ctk.CTkEntry(frame, width=220) + self.capture_rate_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) + self.capture_rate_entry.insert(0, "0.5") + + # blur removal + components.label(frame, 6, 0, "Blur Removal", + tooltip="Threshold for removal of blurry images, relative to all others. \ + For example at 0.2, the blurriest 20%% of the final selected frames will not be saved.") + self.blur_threshold_entry = ctk.CTkEntry(frame, width=220) + self.blur_threshold_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) + self.blur_threshold_entry.insert(0, "0.2") + + # Remove borders + self.image_bordercrop = ctk.BooleanVar(self, False) + components.label(frame, 7, 0, "Remove Borders", + tooltip="Remove black borders from output image") + self.image_bordercrop_entry = ctk.CTkSwitch(frame, variable=self.image_bordercrop, text="") + self.image_bordercrop_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) + + # Crop Variation + components.label(frame, 8, 0, "Crop Variation", + tooltip="Output images will be randomly cropped to +- the base aspect ratio, \ + somewhat biased towards making square images. Set to 0 to use only base sapect.") + self.image_crop_entry = ctk.CTkEntry(frame, width=220) + self.image_crop_entry.grid(row=8, column=1, sticky="w", padx=5, pady=5) + self.image_crop_entry.insert(0, "0.2") + + # # object filter - currently unused, may implement in future + # components.label(frame, 5, 0, "Object Filter", + # tooltip="Detect general features using Haar-Cascade classifier, and choose how to deal with clips where it is detected") + # components.options(frame, 5, 1, ["NONE", "FACE", "EYE", "BODY"], self.video_ui_state, "filter_object") + # components.options(frame, 5, 2, ["INCLUDE", "EXCLUDE", "SUBFOLDER"], self.video_ui_state, "filter_behavior") + + frame.pack(fill="both", expand=1) + return frame + + def __video_download_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0, minsize=120) + frame.grid_columnconfigure(1, weight=0, minsize=200) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # link + components.label(frame, 0, 0, "Single Link", + tooltip="Link to video/playlist to download. Uses yt-dlp, supports youtube, twitch, instagram, and many other sites.") + self.download_link_entry = ctk.CTkEntry(frame, width=220) + self.download_link_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) + components.button(frame, 0, 2, "Download Link", command=lambda: self.__download_button(False)) + + # link list + components.label(frame, 1, 0, "Link List", + tooltip="Path to txt file with list of links separated by newlines.") + self.download_list_entry = ctk.CTkEntry(frame, width=190) + self.download_list_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.download_list_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_file(self.download_list_entry, [("Text file", ".txt")])) + self.download_list_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) + components.button(frame, 1, 2, "Download List", command=lambda: self.__download_button(True)) + + # output directory + components.label(frame, 2, 0, "Output", + tooltip="Path to folder where downloaded videos will be saved.") + self.download_output_entry = ctk.CTkEntry(frame, width=190) + self.download_output_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + self.download_output_button = ctk.CTkButton(frame, width=30, text="...", command=lambda: self.__browse_for_dir(self.download_output_entry)) + self.download_output_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) + + # additional args + 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.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 --format mp4") + components.button(frame, 3, 2, "yt-dlp info", + command=lambda: webbrowser.open("https://github.com/yt-dlp/yt-dlp?tab=readme-ov-file#usage-and-options", new=0, autoraise=False)) + + frame.pack(fill="both", expand=1) + return frame + + def __browse_for_dir(self, entry_box): + # get the path from the user + path = filedialog.askdirectory() + # set the path to the entry box + # delete entry box text + entry_box.focus_set() + entry_box.delete(0, ctk.END) + entry_box.insert(0, path) + self.focus_set() + + def __browse_for_file(self, entry_box, filetypes): + # get the path from the user + path = filedialog.askopenfilename(filetypes=filetypes) + # set the path to the entry box + # delete entry box text + entry_box.focus_set() + entry_box.delete(0, ctk.END) + entry_box.insert(0, path) + self.focus_set() + + def __get_vid_paths(self, batch_mode: bool, input_path_single: str, input_path_dir: str): + input_videos = [] + if not batch_mode: + path = pathlib.Path(input_path_single) + if path.is_file(): + vid = cv2.VideoCapture(str(path)) + ok = False + try: + if vid.isOpened(): + ok, _ = vid.read() + finally: + vid.release() + if ok: + return [path] + else: + self.__update_status("Invalid video file!") + return [] + else: + self.__update_status("No file specified, or invalid file path!") + return [] + else: + input_videos = [] + if not pathlib.Path(input_path_dir).is_dir() or input_path_dir == "": + self.__update_status("Invalid input directory!") + return [] + # Only traverse supported extensions to avoid opening every file. + lower_exts = {e.lower() for e in SUPPORTED_VIDEO_EXTENSIONS} + for path in pathlib.Path(input_path_dir).rglob("*"): + if path.is_file() and path.suffix.lower() in lower_exts: + vid = cv2.VideoCapture(str(path)) + ok = False + try: + if vid.isOpened(): + ok, _ = vid.read() + finally: + vid.release() + if ok: + input_videos.append(path) + self.__update_status(f'Found {len(input_videos)} videos to process') + return input_videos + + def __run_in_thread(self, target, *args): + """Clear status box and run target function in a daemon thread.""" + self.status_label.configure(state="normal") + self.status_label.delete(index1="1.0", index2="end") + self.status_label.configure(state="disabled") + t = threading.Thread(target=target, args=args) + t.daemon = True + t.start() + + @staticmethod + def __parse_timestamp_to_frames(timestamp: str, fps: float) -> int: + return int(sum(int(x) * 60 ** i for i, x in enumerate(reversed(timestamp.split(':')))) * fps) + + def __get_safe_fps(self, video: cv2.VideoCapture, video_path: str) -> float: + fps = video.get(cv2.CAP_PROP_FPS) or 0.0 + if fps <= 0: + self.__update_status(f'Warning: Could not read FPS for "{os.path.basename(video_path)}". Falling back to 30 FPS.') + return 30.0 + return fps + + @staticmethod + def __get_output_dir(use_subdir: bool, batch_mode: bool, output_entry: str, + video_path, input_dir: str) -> str: + if use_subdir and batch_mode: + return os.path.join(output_entry, + os.path.splitext(os.path.relpath(video_path, input_dir))[0]) + elif use_subdir: + return os.path.join(output_entry, + os.path.splitext(os.path.basename(video_path))[0]) + return output_entry + + def __get_random_aspect(self, height: int, width: int, variation: float) -> tuple[int, int, int, int]: + # Return original dimensions and no offset if variation is zero + if variation == 0: + return 0, height, 0, width + + old_aspect = height/width + variation_scaled = old_aspect*variation + if old_aspect > 1.2: #tall image + new_aspect = min(4.0, max(1.0, random.triangular(old_aspect-(variation_scaled*1.5), old_aspect+(variation_scaled/2), old_aspect))) + elif old_aspect < 0.85: #wide image + new_aspect = max(0.25, min(1.0, random.triangular(old_aspect-(variation_scaled/2), old_aspect+(variation_scaled*1.5), old_aspect))) + else: #square image + new_aspect = random.triangular(old_aspect-variation_scaled, old_aspect+variation_scaled) + + new_aspect = round(new_aspect, 2) + #keep the height the same if reducing width, and vice versa + if new_aspect > old_aspect: + new_height = int(height) + new_width = int(width*(old_aspect/new_aspect)) + elif new_aspect < old_aspect: + new_height = int(height*(new_aspect/old_aspect)) + new_width = int(width) + else: + new_height = int(height) + new_width = int(width) + + #random offset in dimension that was cropped + position_x = random.randint(0, width-new_width) + position_y = random.randint(0, height-new_height) + return position_y, new_height, position_x, new_width + + def find_main_contour(self, frame): + #outline image to find main content and exclude black bars often present on letterboxed videos + frame_grayscale = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) + _, frame_thresh = cv2.threshold(frame_grayscale, 15, 255, cv2.THRESH_BINARY) + frame_contours, _ = cv2.findContours(frame_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + if frame_contours: + #select largest contour by area + frame_maincontour = max(frame_contours, key=lambda c: cv2.contourArea(c)) + x1, y1, w1, h1 = cv2.boundingRect(frame_maincontour) + else: #fallback if no contours detected + x1 = 0 + y1 = 0 + h1, w1, _ = frame.shape + + #if bounding box did not detect the correct area, likely due to all-black frame + if not frame_contours or h1 < 10 or w1 < 10: + x1 = 0 + y1 = 0 + h1, w1, _ = frame.shape + return x1, y1, w1, h1 + + def __extract_clips_button(self, batch_mode: bool): + self.__run_in_thread(self.__extract_clips_multi, batch_mode) + + def __extract_clips_multi(self, batch_mode: bool): + if not pathlib.Path(self.clip_output_entry.get()).is_dir() or self.clip_output_entry.get() == "": + self.__update_status("Invalid output directory!") + return + + # validate numeric inputs + try: + max_length = float(self.clip_length_entry.get()) + crop_variation = float(self.clip_crop_entry.get()) + target_fps = float(self.clip_fps_entry.get()) + input_single_entry = self.clip_single_entry.get() + input_multiple_entry = self.clip_list_entry.get() + output_entry = self.clip_output_entry.get() + except ValueError: + self.__update_status("Invalid numeric input for Max Length, Crop Variation, or FPS.") + return + if max_length <= 0.25: + self.__update_status("Max Length of clips must be > 0.25 seconds.") + return + if target_fps < 0: + self.__update_status("Target FPS must be a positive number (or 0 to skip fps re-encoding).") + return + if not (0.0 <= crop_variation < 1.0): + self.__update_status("Crop Variation must be between 0.0 and 1.0.") + return + + input_videos = self.__get_vid_paths(batch_mode, input_single_entry, input_multiple_entry) + if len(input_videos) == 0: # exit if no paths found + return + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + for video_path in input_videos: + output_directory = self.__get_output_dir( + self.output_subdir_clip_entry.get(), batch_mode, + output_entry, video_path, input_multiple_entry) + time_start = "00:00:00" if batch_mode else str(self.clip_time_start_entry.get()) + time_end = "99:99:99" if batch_mode else str(self.clip_time_end_entry.get()) + executor.submit(self.__extract_clips, + str(video_path), time_start, time_end, max_length, + self.split_at_cuts.get(), bool(self.clip_bordercrop_entry.get()), + crop_variation, target_fps, output_directory) + + if batch_mode: + self.__update_status(f'Clip extraction from all videos in "{input_multiple_entry}" complete') + else: + self.__update_status(f'Clip extraction from "{input_single_entry}" complete') + + def __extract_clips(self, video_path: str, timestamp_min: str, timestamp_max: str, max_length: float, + split_at_cuts: bool, remove_borders: bool, crop_variation: float, target_fps: float, output_dir: str): + video = cv2.VideoCapture(video_path) + vid_fps = self.__get_safe_fps(video, video_path) + max_length_frames = int(max_length * vid_fps) + min_length_frames = max(int(0.25 * vid_fps), 1) + total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) or 0 + timestamp_max_frame = min(self.__parse_timestamp_to_frames(timestamp_max, vid_fps), max(total_frames - 1, 0)) + timestamp_min_frame = min(self.__parse_timestamp_to_frames(timestamp_min, vid_fps), timestamp_max_frame) + + if split_at_cuts: + #use scenedetect to find cuts, based on start/end frame number + self.__update_status(f'Detecting scenes in "{os.path.basename(video_path)}"') + timecode_list = scenedetect.detect( + video_path=str(video_path), + detector=scenedetect.AdaptiveDetector(), + start_time=int(timestamp_min_frame), + end_time=int(timestamp_max_frame)) + scene_list = [(x[0].get_frames(), x[1].get_frames()) for x in timecode_list] + if not scene_list: + scene_list = [(timestamp_min_frame, timestamp_max_frame)] + else: + scene_list = [(timestamp_min_frame, timestamp_max_frame)] + + scene_list_split = [] + for scene in scene_list: + length = scene[1]-scene[0] + if length > max_length_frames: #check for any scenes longer than max length + n = math.ceil(length/max_length_frames) #divide into n new scenes + new_length = int(length/n) + new_splits = range(scene[0], scene[1]+min_length_frames, new_length) #divide clip into closest chunks to max_length + for i, _n in enumerate(new_splits[:-1]): + if new_splits[i + 1] - new_splits[i] > min_length_frames: + scene_list_split.append((new_splits[i], new_splits[i + 1])) + elif length > (min_length_frames + 2): + # Trim first/last frame to avoid transition artifacts + scene_list_split.append((scene[0] + 1, scene[1] - 1)) + + self.__update_status(f'Video "{os.path.basename(video_path)}" being split into {len(scene_list_split)} clips in "{output_dir}"') + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + futures = [ + executor.submit(self.__save_clip, scene, video_path, target_fps, + remove_borders, crop_variation, output_dir) + for scene in scene_list_split + ] + for future in concurrent.futures.as_completed(futures): + exc = future.exception() + if exc is not None: + self.__update_status(f'Error saving clip: {exc}') + + video.release() + + def __save_clip(self, scene: tuple[int, int], video_path: str, target_fps: float, + remove_borders: bool, crop_variation: float, output_dir: str): + basename, ext = os.path.splitext(os.path.basename(video_path)) + video = cv2.VideoCapture(str(video_path)) + fps = self.__get_safe_fps(video, video_path) + os.makedirs(output_dir, exist_ok=True) + output_name = f'{output_dir}{os.sep}{basename}_{scene[0]}-{scene[1]}' + output_ext = ".mp4" + + video.set(cv2.CAP_PROP_POS_FRAMES, (scene[1] + scene[0])//2) #set to middle of scene + frame_number = int(video.get(cv2.CAP_PROP_POS_FRAMES)) + success, frame = video.read() + if not success or frame is None: + self.__update_status(f'Failed to read frame from "{os.path.basename(video_path)}" at {int(frame_number)}. Skipping clip.') + video.release() + return + + # Blend random frames to detect borders, avoiding incorrect crop from black frames + if remove_borders: + frame_blend = frame + for i in range(5): + random_frame = random.randint(scene[0], scene[1]) + video.set(cv2.CAP_PROP_POS_FRAMES, random_frame) + success, frame = video.read() + if not success or frame is None: + continue + a = 1/(i+1) + b = 1-a + frame_blend = cv2.addWeighted(frame, a, frame_blend, b, 0) + x1, y1, w1, h1 = self.find_main_contour(frame_blend) + else: + x1 = 0 + y1 = 0 + h1, w1, _ = frame.shape + + y2, h2, x2, w2 = self.__get_random_aspect(h1, w1, crop_variation) + # Ensure dimensions are even, required + h2 -= h2 % 2 + w2 -= w2 % 2 + print(end='\x1b[2K') #clear terminal so next line can overwrite it + print(f'Saving frames {scene[0]}-{scene[1]} at size {w2}x{h2}', end="\r") + video.set(cv2.CAP_PROP_POS_FRAMES, (scene[1] + scene[0])//2) + success, frame = video.read() + if success: + try: + preview = Image.fromarray( + cv2.cvtColor(frame[y1+y2:y1+y2+h2, x1+x2:x1+x2+w2], cv2.COLOR_BGR2RGB)) + preview.thumbnail((150, 150)) + self.preview_image.configure(light_image=preview, size=preview.size) + #truncate filename of long files so UI doesn't shift around + filename_truncated = basename + ext if len(basename) < 20 else basename[:18] + ".." + ext + self.preview_image_label.configure( + text=f'{filename_truncated}\nFrames: {scene[0]}-{scene[1]}\nSize: {w2}x{h2}') + except Exception: + pass + video.release() + + if target_fps <= 0: + target_fps = fps + + output_path = f'{output_name}{output_ext}' + self.__write_clip_av(video_path, output_path, scene, fps, target_fps, + x1 + x2, y1 + y2, w2, h2) + + @staticmethod + def __write_clip_av(video_path: str, output_path: str, scene: tuple[int, int], + src_fps: float, target_fps: float, + crop_x: int, crop_y: int, crop_w: int, crop_h: int): + start_sec = scene[0] / src_fps + end_sec = scene[1] / src_fps + rate_frac = Fraction(target_fps).limit_denominator(10000) + stream_time_base = Fraction(rate_frac.denominator, rate_frac.numerator) + + with av.open(video_path) as input_container: + in_video = input_container.streams.video[0] + in_video.thread_type = 'AUTO' + in_audio = input_container.streams.audio[0] if input_container.streams.audio else None + + with av.open(output_path, mode='w') as output_container: + out_video = output_container.add_stream('libx264', rate=rate_frac) + out_video.width = crop_w + out_video.height = crop_h + out_video.pix_fmt = 'yuv420p' + out_video.time_base = stream_time_base + + out_audio = output_container.add_stream_from_template(in_audio) if in_audio else None + + input_container.seek(int(start_sec * 1_000_000)) + + out_frame_idx = 0 + out_time_step = 1.0 / target_fps + video_done = False + decode_streams = [s for s in (in_video, in_audio) if s is not None] + + for packet in input_container.demux(decode_streams): + if packet.stream == in_video: + if video_done: + continue + for frame in packet.decode(): + if frame.time is None or frame.time < start_sec: + continue + if frame.time >= end_sec: + video_done = True + break + + # FPS conversion: skip frames when source fps > target fps + if frame.time < start_sec + out_frame_idx * out_time_step: + continue + + img = frame.to_ndarray(format='bgr24') + cropped = img[crop_y:crop_y + crop_h, crop_x:crop_x + crop_w] + out_frame = av.VideoFrame.from_ndarray(cropped, format='bgr24') + out_frame.pts = out_frame_idx + out_frame.time_base = stream_time_base + + for out_pkt in out_video.encode(out_frame): + output_container.mux(out_pkt) + out_frame_idx += 1 + + elif packet.stream == in_audio and out_audio is not None: + if packet.dts is None: + continue + pkt_time = float(packet.pts * packet.time_base) + if pkt_time < start_sec or pkt_time >= end_sec: + continue + # Re-timestamp audio relative to clip start + packet.pts = int((pkt_time - start_sec) / packet.time_base) + packet.dts = packet.pts + packet.stream = out_audio + output_container.mux(packet) + + # Flush video encoder + for pkt in out_video.encode(): + output_container.mux(pkt) + + def __extract_images_button(self, batch_mode: bool): + self.__run_in_thread(self.__extract_images_multi, batch_mode) + + def __extract_images_multi(self, batch_mode : bool): + if not pathlib.Path(self.image_output_entry.get()).is_dir() or self.image_output_entry.get() == "": + self.__update_status("Invalid output directory!") + return + + # validate numeric inputs + try: + capture_rate = float(self.capture_rate_entry.get()) + blur_threshold = float(self.blur_threshold_entry.get()) + crop_variation = float(self.image_crop_entry.get()) + input_single_entry = self.image_single_entry.get() + input_multiple_entry = self.image_list_entry.get() + output_entry = self.image_output_entry.get() + except ValueError: + self.__update_status("Invalid numeric input for Images/sec, Blur Removal, or Crop Variation.") + return + if capture_rate <= 0: + self.__update_status("Images/sec must be > 0.") + return + if not (0.0 <= blur_threshold < 1.0): + self.__update_status("Blur Removal must be between 0.0 and 1.0.") + return + if not (0.0 <= crop_variation < 1.0): + self.__update_status("Crop Variation must be between 0.0 and 1.0.") + return + + input_videos = self.__get_vid_paths(batch_mode, input_single_entry, input_multiple_entry) + if not input_videos: + return + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + for video_path in input_videos: + output_directory = self.__get_output_dir( + self.output_subdir_img_entry.get(), batch_mode, + output_entry, video_path, input_multiple_entry) + time_start = "00:00:00" if batch_mode else str(self.image_time_start_entry.get()) + time_end = "99:99:99" if batch_mode else str(self.image_time_end_entry.get()) + executor.submit(self.__save_frames, + str(video_path), time_start, time_end, capture_rate, + blur_threshold, self.image_bordercrop.get(), + crop_variation, output_directory) + if batch_mode: + self.__update_status(f'Image extraction from all videos in {input_multiple_entry} complete') + else: + self.__update_status(f'Image extraction from "{input_single_entry}" complete') + + def __save_frames(self, video_path: str, timestamp_min: str, timestamp_max: str, capture_rate: float, + blur_threshold: float, remove_borders: bool, crop_variation: float, output_dir: str): + video = cv2.VideoCapture(video_path) + vid_fps = self.__get_safe_fps(video, video_path) + if capture_rate <= 0: + self.__update_status("Images/sec must be > 0.") + video.release() + return + image_rate = max(int(vid_fps / capture_rate), 1) + total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) or 0 + timestamp_max_frame = min(self.__parse_timestamp_to_frames(timestamp_max, vid_fps), max(total_frames - 1, 0)) + timestamp_min_frame = min(self.__parse_timestamp_to_frames(timestamp_min, vid_fps), timestamp_max_frame) + frame_range = range(timestamp_min_frame, timestamp_max_frame, image_rate) + frame_list = [] + + for n in frame_range: + #pick frame from random triangular distribution around center of each "chunk" of the video + frame = abs(int(random.triangular(n-(image_rate/2), n+(image_rate/2)))) + frame = max(0, min(frame, max(total_frames - 1, 0))) + frame_list.append(frame) + + self.__update_status(f'Video "{os.path.basename(video_path)}" will be split into {len(frame_list)} images in "{output_dir}"') + + output_list = [] + for f in frame_list: + video.set(cv2.CAP_PROP_POS_FRAMES, f) + success, frame = video.read() + if success and frame is not None: + frame_grayscale = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) + frame_sharpness = cv2.Laplacian(frame_grayscale, cv2.CV_64F).var() + output_list.append((f, frame_sharpness)) + + if not output_list: + self.__update_status(f'No frames extracted from "{os.path.basename(video_path)}" in the selected range.') + video.release() + return + + output_list_sorted = sorted(output_list, key=lambda x: x[1]) + cutoff = int(blur_threshold * len(output_list_sorted)) + output_list_cut = output_list_sorted[cutoff:] + self.__update_status(f'{cutoff} blurriest images have been dropped from "{os.path.basename(video_path)}"') + + basename, ext = os.path.splitext(os.path.basename(video_path)) + os.makedirs(output_dir, exist_ok=True) + + for f in output_list_cut: + filename = f'{output_dir}{os.sep}{basename}_{f[0]}.jpg' + video.set(cv2.CAP_PROP_POS_FRAMES, f[0]) + success, frame = video.read() + + #crop out borders of frame + if remove_borders and success and frame is not None: + x1, y1, w1, h1 = self.find_main_contour(frame) + frame_cropped = frame[y1:y1+h1, x1:x1+w1] + else: + frame_cropped = frame if success and frame is not None else None + if frame_cropped is not None: + x1 = 0 + y1 = 0 + h1, w1, _ = frame_cropped.shape + + y2, h2, x2, w2 = self.__get_random_aspect(h1, w1, crop_variation) + + if success and frame is not None and frame_cropped is not None: + print(end='\x1b[2K') #clear terminal so next line can overwrite it + print(f'Saving frame {f[0]} at size {w2}x{h2}', end="\r") + try: + preview = Image.fromarray( + cv2.cvtColor(frame_cropped[y2:y2+h2, x2:x2+w2], cv2.COLOR_BGR2RGB)) + preview.thumbnail((150, 150)) + filename_truncated = basename + ext if len(basename) < 20 else basename[:17] + "..." + ext + self.preview_image.configure(light_image=preview, size=preview.size) + self.preview_image_label.configure(text=f'{filename_truncated}\nFrame: {f[0]}\nSize: {w2}x{h2}') + except Exception: + pass # preview update is non-critical + + cv2.imwrite(filename, frame_cropped[y2:y2+h2, x2:x2+w2]) + video.release() + + def __download_button(self, batch_mode: bool): + self.__run_in_thread(self.__download_multi, batch_mode) + + def __update_status(self, status_text: str): + print(status_text) + self.status_label.configure(state="normal") + self.status_label.insert(index="end", text=status_text + "\n") + self.status_label.configure(state="disabled") + + def __download_multi(self, batch_mode: bool): + if not pathlib.Path(self.download_output_entry.get()).is_dir() or self.download_output_entry.get() == "": + self.__update_status("Invalid output directory!") + return + + if not batch_mode: + ydl_urls = [self.download_link_entry.get()] + elif batch_mode: + ydl_path = pathlib.Path(self.download_list_entry.get()) + if ydl_path.is_file() and ydl_path.suffix.lower() == ".txt": + with open(ydl_path) as file: + ydl_urls = file.readlines() + else: + self.__update_status("Invalid link list!") + return + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + for url in ydl_urls: + executor.submit(self.__download_video, + url.strip(), self.download_output_entry.get(), + self.download_args_entry.get("0.0", ctk.END)) + + self.__update_status(f'Completed {len(ydl_urls)} downloads.') + + def __download_video(self, url: str, output_dir: str, output_args: str): + url = (url or "").strip() + if not url: + self.__update_status("Empty URL, skipping download.") + return + + #Respect quotes and split into list to run as yt-dlp command + additional_args = shlex.split(output_args.strip()) if output_args and output_args.strip() else [] + cmd = ["yt-dlp", "-o", "%(title)s.%(ext)s", "-P", output_dir] + additional_args + [url] + + self.__update_status(f'Downloading {url}') + subprocess.run(cmd) + self.__update_status(f'Download {url} done!') diff --git a/modules/util/ui/ctk_validation.py b/modules/util/ui/ctk_validation.py new file mode 100644 index 000000000..9f44de8eb --- /dev/null +++ b/modules/util/ui/ctk_validation.py @@ -0,0 +1,501 @@ +from __future__ import annotations + +import contextlib +import os +import re +import sys +import tkinter as tk +from collections import deque +from collections.abc import Callable +from pathlib import PurePosixPath, PureWindowsPath +from typing import TYPE_CHECKING, Any +from urllib.parse import urlparse + +from modules.util.enum.ModelFormat import ModelFormat +from modules.util.enum.PathIOType import PathIOType + +if TYPE_CHECKING: + from modules.util.ui.UIState import UIState + + import customtkinter as ctk + + +DEBOUNCE_TYPING_MS = 250 +UNDO_DEBOUNCE_MS = 500 +ERROR_BORDER_COLOR = "#dc3545" + +_active_validators: set[FieldValidator] = set() + +TRAILING_SLASH_RE = re.compile(r"[\\/]$") +ENDS_WITH_EXT = re.compile(r"\.[A-Za-z0-9]+$") +HUGGINGFACE_REPO_RE = re.compile(r"^[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+$") + +_INVALID_CHARS = {chr(c) for c in range(32)} +_IS_WINDOWS = sys.platform == "win32" +if _IS_WINDOWS: + _INVALID_CHARS |= set('<>"|?*') + + +def _is_huggingface_repo_or_file(value: str) -> bool: + trimmed = value.strip() + + if trimmed.startswith("https://"): + parsed = urlparse(trimmed) + if parsed.netloc not in {"huggingface.co", "huggingface.com"}: + return False + parts = parsed.path.strip("/").split("/") + if len(parts) >= 5 and parts[2] in {"resolve", "blob"}: + return bool(ENDS_WITH_EXT.search(parts[-1])) + return False + + if len(trimmed) > 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 len(trimmed) >= 2 and trimmed[1] == ":" and trimmed[0].isalpha(): + return False + if trimmed.count("/") != 1: + return False + + return bool(HUGGINGFACE_REPO_RE.match(trimmed)) + + +def _has_invalid_chars(value: str) -> bool: + return bool(_INVALID_CHARS.intersection(value)) + + +def _check_overwrite(path: str, *, is_dir: bool, prevent: bool) -> str | None: + if not prevent: + return None + abs_path = os.path.abspath(path) + if is_dir and os.path.isdir(abs_path): + return "Output folder already exists (overwrite prevented)" + if not is_dir and os.path.isfile(abs_path): + return "Output file already exists (overwrite prevented)" + return None + + +def validate_path( + value: str, + io_type: PathIOType = PathIOType.INPUT, + *, + prevent_overwrites: bool = False, + output_format: str | None = None, +) -> str | None: + """Return an error string if *value* is an invalid path, else ``None``.""" + trimmed = value.strip() + + if not trimmed: + return "Path is empty" + if TRAILING_SLASH_RE.search(trimmed): + return "Path must not end with a slash" + if _has_invalid_chars(trimmed): + return "Path contains invalid characters" + + if trimmed.startswith("cloud:"): + cloud_path = trimmed[6:] + if not cloud_path: + return "Cloud path is empty" + if cloud_path.startswith(("http://", "https://")): + return "Cloud path cannot be a URL" + if "\\" in cloud_path: + return "Cloud path must use forward slashes (/)" + return None + + if io_type == PathIOType.INPUT and _is_huggingface_repo_or_file(trimmed): + return None + + if io_type == PathIOType.INPUT: + if not os.path.exists(os.path.abspath(trimmed)): + return "Input path does not exist" + + if io_type in (PathIOType.OUTPUT, PathIOType.MODEL): + if not os.path.isdir(os.path.dirname(os.path.abspath(trimmed))): + return "Parent folder does not exist" + + if io_type == PathIOType.MODEL and output_format is not None: + if output_format == "DIFFUSERS": + if ENDS_WITH_EXT.search(trimmed): + return "Diffusers output must be a directory path, not a file" + return _check_overwrite(trimmed, is_dir=True, prevent=prevent_overwrites) + + try: + expected_ext = ModelFormat[output_format].file_extension() + except KeyError: + expected_ext = "" + + if expected_ext: + suffix = (PureWindowsPath(trimmed) if _IS_WINDOWS else PurePosixPath(trimmed)).suffix.lower() + if suffix != expected_ext: + return f"Extension must be '{expected_ext}' for {output_format} format" + return _check_overwrite(trimmed, is_dir=False, prevent=prevent_overwrites) + + if io_type == PathIOType.OUTPUT: + return _check_overwrite(trimmed, is_dir=False, prevent=prevent_overwrites) + + return None + +DEFAULT_MAX_UNDO = 20 + + +class UndoHistory: + def __init__(self, max_size: int = DEFAULT_MAX_UNDO): + self._stack: deque[str] = deque(maxlen=max_size) + self._redo_stack: list[str] = [] + + def push(self, value: str): + if self._stack and self._stack[-1] == value: + return + self._stack.append(value) + self._redo_stack.clear() + + def undo(self, current: str) -> str | None: + if not self._stack: + return None + top = self._stack[-1] + if top == current and len(self._stack) > 1: + self._redo_stack.append(self._stack.pop()) + return self._stack[-1] + elif top != current: + self._redo_stack.append(current) + return top + return None + + def redo(self) -> str | None: + if not self._redo_stack: + return None + value = self._redo_stack.pop() + self._stack.append(value) + return value + + +class DebounceTimer: + def __init__(self, widget, delay_ms: int, callback: Callable[..., Any]): + self.widget = widget + self.delay_ms = delay_ms + self.callback = callback + self._after_id: str | None = None + + def call(self, *args, **kwargs): + if self._after_id: + with contextlib.suppress(tk.TclError): + self.widget.after_cancel(self._after_id) + + def fire(): + self._after_id = None + self.callback(*args, **kwargs) + + with contextlib.suppress(tk.TclError): + self._after_id = self.widget.after(self.delay_ms, fire) + + def cancel(self): + if self._after_id: + with contextlib.suppress(tk.TclError): + self.widget.after_cancel(self._after_id) + self._after_id = None + + +class FieldValidator: + def __init__( + self, + component: ctk.CTkEntry, + var: tk.Variable, + ui_state: UIState, + var_name: str, + max_undo: int = DEFAULT_MAX_UNDO, + extra_validate: Callable[[str], str | None] | None = None, + required: bool = False, + ): + self.component = component + self.var = var + self.ui_state = ui_state + self.var_name = var_name + self._extra_validate = extra_validate + self._required = required + + try: + self._original_border_color = component.cget("border_color") + except Exception: + self._original_border_color = "gray50" + + self._shadow_var = tk.StringVar(master=component) + self._shadow_trace_name: str | None = None + self._real_var_trace_name: str | None = None + self._syncing = False + self._touched = False + self._bound = False + + self._debounce: DebounceTimer | None = None + self._undo_debounce: DebounceTimer | None = None + self._undo = UndoHistory(max_undo) + + def attach(self) -> None: + self._shadow_var.set(self.var.get()) + self._swap_textvariable(self._shadow_var) + + self._debounce = DebounceTimer( + self.component, DEBOUNCE_TYPING_MS, self._on_debounce_fire + ) + self._undo_debounce = DebounceTimer( + self.component, UNDO_DEBOUNCE_MS, self._push_undo_snapshot + ) + + self._shadow_trace_name = self._shadow_var.trace_add("write", self._on_shadow_write) + self._real_var_trace_name = self.var.trace_add("write", self._on_real_var_write) + + self.component.bind("", self._on_focus_in) + self.component.bind("", self._on_user_input) + self.component.bind("<>", self._on_user_input) + self.component.bind("<>", self._on_user_input) + self.component.bind("", self._on_focus_out) + self.component.bind("", self._on_undo) + self.component.bind("", self._on_undo) + self.component.bind("", self._on_redo) + self.component.bind("", self._on_redo) + self.component.bind("", self._on_redo) + self.component.bind("", self._on_redo) + self.component.bind("", self._on_enter) + + self._bound = True + _active_validators.add(self) + + def detach(self) -> None: + if not self._bound: + return + self._bound = False + _active_validators.discard(self) + + self._commit() + + if self._debounce: + self._debounce.cancel() + if self._undo_debounce: + self._undo_debounce.cancel() + + if self._shadow_trace_name: + with contextlib.suppress(Exception): + self._shadow_var.trace_remove("write", self._shadow_trace_name) + self._shadow_trace_name = None + + if self._real_var_trace_name: + with contextlib.suppress(Exception): + self.var.trace_remove("write", self._real_var_trace_name) + self._real_var_trace_name = None + + self._swap_textvariable(self.var) + + def _swap_textvariable(self, new_var: tk.Variable) -> None: + comp = self.component + if comp._textvariable_callback_name: + with contextlib.suppress(Exception): + comp._textvariable.trace_remove("write", comp._textvariable_callback_name) # type: ignore[union-attr] + comp._textvariable_callback_name = "" + + comp.configure(textvariable=new_var) + + if new_var is not None: + comp._textvariable_callback_name = new_var.trace_add( + "write", comp._textvariable_callback + ) + + def _commit(self) -> None: + shadow_val = self._shadow_var.get() + if shadow_val != self.var.get(): + self._syncing = True + self.var.set(shadow_val) + self._syncing = False + + def validate(self, value: str) -> str | None: + """Return an error string if *value* is invalid, else None.""" + meta = self.ui_state.get_field_metadata(self.var_name) + declared_type = meta.type + nullable = meta.nullable + default_val = meta.default + + if value == "": + if self._required: + return "Value required" + if nullable: + return None + if declared_type is str: + if default_val == "": + return None + return "Value required" + return None + + try: + if declared_type is int: + v = int(value) + if v < 0: + return "Value must be non-negative" + elif declared_type is float: + v = float(value) + if v < 0: + return "Value must be non-negative" + elif declared_type is bool: + if value.lower() not in ("true", "false", "0", "1"): + return "Invalid bool" + except ValueError: + return "Invalid value" + + if self._extra_validate is not None: + return self._extra_validate(value) + + return None + + def _apply_error(self) -> None: + self.component.configure(border_color=ERROR_BORDER_COLOR) + + def _clear_error(self) -> None: + self.component.configure(border_color=self._original_border_color) + + def _validate_and_style(self, value: str) -> bool: + error = self.validate(value) + if error is None: + self._clear_error() + return True + else: + self._apply_error() + return False + + def _on_shadow_write(self, *_args) -> None: + if self._syncing: + return + if not self._touched: + # external sync or initial set — commit immediately + self._commit() + if self._debounce: + self._debounce.cancel() + return + if self._debounce: + self._debounce.call() + if self._undo_debounce: + self._undo_debounce.call() + + def _on_real_var_write(self, *_args) -> None: + if self._syncing: + return + # external change (preset load, file dialog, etc) — sync to shadow var + self._syncing = True + self._shadow_var.set(self.var.get()) + self._syncing = False + self._validate_and_style(self._shadow_var.get()) + + def _push_undo_snapshot(self) -> None: + self._undo.push(self._shadow_var.get()) + + def _on_debounce_fire(self) -> None: + val = self._shadow_var.get() + if self._validate_and_style(val): + self._commit() + + def _on_focus_in(self, _e=None) -> None: + self._touched = False + self._undo.push(self._shadow_var.get()) + + def _on_user_input(self, _e=None) -> None: + self._touched = True + + def _on_focus_out(self, _e=None) -> None: + if self._debounce: + self._debounce.cancel() + if self._undo_debounce: + self._undo_debounce.cancel() + if self._touched: + if self._validate_and_style(self._shadow_var.get()): + self._commit() + self._undo.push(self._shadow_var.get()) + + def _on_enter(self, _e=None) -> None: + if self._debounce: + self._debounce.cancel() + if self._touched: + if self._validate_and_style(self._shadow_var.get()): + self._commit() + + def _set_value(self, value: str) -> None: + self._syncing = True + self._shadow_var.set(value) + self._syncing = False + if self._validate_and_style(value): + self._commit() + + def _on_undo(self, _e=None) -> str: + previous = self._undo.undo(self._shadow_var.get()) + if previous is not None: + self._set_value(previous) + return "break" + + def _on_redo(self, _e=None) -> str: + next_val = self._undo.redo() + if next_val is not None: + self._set_value(next_val) + return "break" + + +class PathValidator(FieldValidator): + """FieldValidator with additional path-specific checks.""" + + def __init__( + self, + component: ctk.CTkEntry, + var: tk.Variable, + ui_state: UIState, + var_name: str, + io_type: PathIOType = PathIOType.INPUT, + max_undo: int = DEFAULT_MAX_UNDO, + extra_validate: Callable[[str], str | None] | None = None, + required: bool = False, + ): + super().__init__(component, var, ui_state, var_name, max_undo=max_undo, extra_validate=extra_validate, required=required) + self.io_type = io_type + + def _get_var_safe(self, name: str) -> tk.Variable | None: + try: + return self.ui_state.get_var(name) + except (KeyError, AttributeError): + return None + + def validate(self, value: str) -> str | None: + base_err = super().validate(value) + if base_err is not None: + return base_err + if value == "": + return None + + prevent_var = self._get_var_safe("prevent_overwrites") + format_var = self._get_var_safe("output_model_format") + return validate_path( + value, + io_type=self.io_type, + prevent_overwrites=prevent_var.get() if prevent_var is not None else False, + output_format=format_var.get() if format_var is not None else None, + ) + + def revalidate(self) -> None: + if self.component.winfo_exists(): + self._validate_and_style(self._shadow_var.get()) + + +def flush_and_validate_all() -> list[str]: + invalid: list[str] = [] + + for v in list(_active_validators): + if v._debounce: + v._debounce.cancel() + + value = v._shadow_var.get() + error = v.validate(value) + + if error is not None: + v._apply_error() + invalid.append(f"{v.var_name}: {error}") + else: + v._clear_error() + v._commit() + + return invalid From a88f369790c73bcae356382165a89b90c31f7a63 Mon Sep 17 00:00:00 2001 From: dxqb Date: Sun, 10 May 2026 14:06:56 +0200 Subject: [PATCH 02/21] refactor: abstract CTK views and migrate logic to controllers Co-Authored-By: Claude Sonnet 4.6 --- .gitignore | 1 + .../ui/AdditionalEmbeddingsTabController.py | 14 + modules/ui/BaseAdditionalEmbeddingsTabView.py | 127 +- modules/ui/BaseCaptionUIView.py | 570 +------- modules/ui/BaseCloudTabView.py | 233 ++-- modules/ui/BaseConceptTabView.py | 223 +--- modules/ui/BaseConceptWindowView.py | 985 ++++---------- modules/ui/BaseConfigListView.py | 163 +-- modules/ui/BaseConvertModelUIView.py | 137 +- modules/ui/BaseGenerateCaptionsWindowView.py | 134 +- modules/ui/BaseGenerateMasksWindowView.py | 152 +-- modules/ui/BaseLoraTabView.py | 168 +-- modules/ui/BaseModelTabView.py | 476 +++---- modules/ui/BaseMuonAdamWindowView.py | 73 +- modules/ui/BaseOffloadingWindowView.py | 87 +- modules/ui/BaseOptimizerParamsWindowView.py | 171 +-- modules/ui/BaseProfilingWindowView.py | 66 +- modules/ui/BaseSampleFrameView.py | 137 +- modules/ui/BaseSampleParamsWindowView.py | 39 +- modules/ui/BaseSampleWindowView.py | 225 +--- modules/ui/BaseSamplingTabView.py | 121 +- modules/ui/BaseSchedulerParamsWindowView.py | 122 +- .../ui/BaseTimestepDistributionWindowView.py | 176 +-- modules/ui/BaseTopBarView.py | 186 +-- modules/ui/BaseTrainUIView.py | 919 +++---------- modules/ui/BaseTrainingTabView.py | 1148 ++++++++--------- modules/ui/BaseVideoToolUIView.py | 905 ++----------- modules/ui/CaptionUIController.py | 369 ++---- modules/ui/CloudTabController.py | 31 + modules/ui/ConceptTabController.py | 19 + modules/ui/ConceptWindowController.py | 723 +---------- modules/ui/ConvertModelUIController.py | 116 +- modules/ui/CtkAdditionalEmbeddingsTabView.py | 119 +- modules/ui/CtkCaptionUIView.py | 476 +------ modules/ui/CtkCloudTabView.py | 210 +-- modules/ui/CtkConceptTabView.py | 196 +-- modules/ui/CtkConceptWindowView.py | 903 +------------ modules/ui/CtkConfigListView.py | 369 +----- modules/ui/CtkConvertModelUIView.py | 160 +-- modules/ui/CtkGenerateCaptionsWindowView.py | 41 +- modules/ui/CtkGenerateMasksWindowView.py | 44 +- modules/ui/CtkLoraTabView.py | 131 +- modules/ui/CtkModelTabView.py | 674 +--------- modules/ui/CtkMuonAdamWindowView.py | 104 +- modules/ui/CtkOffloadingWindowView.py | 70 +- modules/ui/CtkOptimizerParamsWindowView.py | 259 +--- modules/ui/CtkProfilingWindowView.py | 54 +- modules/ui/CtkSampleFrameView.py | 111 +- modules/ui/CtkSampleParamsWindowView.py | 29 +- modules/ui/CtkSampleWindowView.py | 167 +-- modules/ui/CtkSamplingTabView.py | 116 +- modules/ui/CtkSchedulerParamsWindowView.py | 77 +- .../ui/CtkTimestepDistributionWindowView.py | 146 +-- modules/ui/CtkTopBarView.py | 252 +--- modules/ui/CtkTrainUIView.py | 778 +++-------- modules/ui/CtkTrainingTabView.py | 847 +----------- modules/ui/CtkVideoToolUIView.py | 911 ++----------- .../ui/GenerateCaptionsWindowController.py | 135 +- modules/ui/GenerateMasksWindowController.py | 153 +-- modules/ui/LoraTabController.py | 22 + modules/ui/ModelTabController.py | 13 + modules/ui/MuonAdamWindowController.py | 28 + modules/ui/OffloadingWindowController.py | 6 + modules/ui/OptimizerParamsWindowController.py | 273 +--- modules/ui/ProfilingWindowController.py | 25 + modules/ui/SampleFrameController.py | 17 + modules/ui/SampleParamsWindowController.py | 8 + modules/ui/SampleWindowController.py | 110 +- modules/ui/SamplingTabController.py | 14 + modules/ui/SchedulerParamsWindowController.py | 17 + .../TimestepDistributionWindowController.py | 146 +-- modules/ui/TopBarController.py | 281 +--- modules/ui/TrainUIController.py | 885 ++----------- modules/ui/TrainingTabController.py | 39 + modules/ui/VideoToolUIController.py | 452 ++----- modules/util/path_util.py | 6 + modules/util/ui/CtkUIState.py | 23 + modules/util/ui/UIState.py | 86 +- .../ui/{components.py => ctk_components.py} | 75 +- modules/util/ui/ctk_validation.py | 276 +--- modules/util/ui/validation.py | 300 +---- scripts/train_ui.py | 4 +- 82 files changed, 3652 insertions(+), 16002 deletions(-) create mode 100644 modules/ui/AdditionalEmbeddingsTabController.py create mode 100644 modules/ui/CloudTabController.py create mode 100644 modules/ui/ConceptTabController.py create mode 100644 modules/ui/LoraTabController.py create mode 100644 modules/ui/ModelTabController.py create mode 100644 modules/ui/MuonAdamWindowController.py create mode 100644 modules/ui/OffloadingWindowController.py create mode 100644 modules/ui/ProfilingWindowController.py create mode 100644 modules/ui/SampleFrameController.py create mode 100644 modules/ui/SampleParamsWindowController.py create mode 100644 modules/ui/SamplingTabController.py create mode 100644 modules/ui/SchedulerParamsWindowController.py create mode 100644 modules/ui/TrainingTabController.py create mode 100644 modules/util/ui/CtkUIState.py rename modules/util/ui/{components.py => ctk_components.py} (88%) diff --git a/.gitignore b/.gitignore index dce73072c..da19fd74e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ # development +PLAN.md # claude's planning file, kept local-only .idea *.bak *.pyc diff --git a/modules/ui/AdditionalEmbeddingsTabController.py b/modules/ui/AdditionalEmbeddingsTabController.py new file mode 100644 index 000000000..638df8e68 --- /dev/null +++ b/modules/ui/AdditionalEmbeddingsTabController.py @@ -0,0 +1,14 @@ + +from modules.util.config.TrainConfig import TrainConfig, TrainEmbeddingConfig + + +class AdditionalEmbeddingsTabController: + def __init__(self, config: TrainConfig): + self.train_config = config + + def create_new_element(self) -> TrainEmbeddingConfig: + return TrainEmbeddingConfig.default_values() + + def randomize_uuid(self, embedding_config: TrainEmbeddingConfig) -> TrainEmbeddingConfig: + embedding_config.uuid = TrainEmbeddingConfig.default_values().uuid + return embedding_config diff --git a/modules/ui/BaseAdditionalEmbeddingsTabView.py b/modules/ui/BaseAdditionalEmbeddingsTabView.py index 6a5e3fbe7..744f0abdc 100644 --- a/modules/ui/BaseAdditionalEmbeddingsTabView.py +++ b/modules/ui/BaseAdditionalEmbeddingsTabView.py @@ -1,136 +1,75 @@ -from modules.ui.ConfigList import ConfigList -from modules.util.config.TrainConfig import TrainConfig, TrainEmbeddingConfig -from modules.util.ui import components -from modules.util.ui.UIState import UIState +from modules.ui.BaseConfigListView import BaseConfigListView +from modules.util import path_util import customtkinter as ctk -class AdditionalEmbeddingsTab(ConfigList): - - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__( - master, - train_config, - ui_state, - attr_name="additional_embeddings", - enable_key="train", - from_external_file=False, - add_button_text="add embedding", - is_full_width=True, - show_toggle_button=True - ) +class BaseAdditionalEmbeddingsTabView(BaseConfigListView): def refresh_ui(self): if self.element_list is not None: - self.element_list.destroy() + self._destroy_frame(self.element_list) self.element_list = None self.widgets_initialized = False self._create_element_list() - def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - return EmbeddingWidget(master, element, i, open_command, remove_command, clone_command, save_command) - - def create_new_element(self) -> dict: - return TrainEmbeddingConfig.default_values() - def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: pass -class EmbeddingWidget(ctk.CTkFrame): - def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): - super().__init__( - master=master, corner_radius=10, bg_color="transparent" - ) +class BaseEmbeddingWidgetView: + + def __init__(self, components): + self.components = components - self.element = element - self.ui_state = UIState(self, element) + def build_content(self, top_frame, bottom_frame, ui_state, i, save_command, remove_command, clone_command, controller): + self.ui_state = ui_state self.i = i self.save_command = save_command - self.grid_columnconfigure(0, weight=1) - - top_frame = ctk.CTkFrame(master=self, corner_radius=0, fg_color="transparent") - top_frame.grid(row=0, column=0, sticky="nsew") - top_frame.grid_columnconfigure(3, weight=1) - top_frame.grid_columnconfigure(5, weight=1) - - bottom_frame = ctk.CTkFrame(master=self, corner_radius=0, fg_color="transparent") - bottom_frame.grid(row=1, column=0, sticky="nsew") - bottom_frame.grid_columnconfigure(7, weight=1) - # close button - close_button = ctk.CTkButton( - master=top_frame, - width=20, - height=20, - text="X", - corner_radius=2, - fg_color="#C00000", - command=lambda: remove_command(self.i), - ) - close_button.grid(row=0, column=0) + self.components.colored_icon_button(top_frame, 0, 0, "X", "#C00000", lambda: remove_command(self.i)) # clone button - clone_button = ctk.CTkButton( - master=top_frame, - width=20, - height=20, - text="+", - corner_radius=2, - fg_color="#00C000", - command=lambda: clone_command(self.i, self.__randomize_uuid), - ) - clone_button.grid(row=0, column=1, padx=5) + self.components.colored_icon_button(top_frame, 0, 1, "+", "#00C000", lambda: clone_command(self.i, controller.randomize_uuid), padx=5) # 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.path_entry( + self.components.label(top_frame, 0, 2, "base embedding:", + tooltip="The base embedding to train on. Leave empty to create a new embedding") + self.components.path_entry( top_frame, 0, 3, self.ui_state, "model_name", - mode="file", path_modifier=components.json_path_modifier + mode="file", path_modifier=path_util.json_path_modifier ) # placeholder - components.label(top_frame, 0, 4, "placeholder:", - tooltip="The placeholder used when using the embedding in a prompt") - components.entry(top_frame, 0, 5, self.ui_state, "placeholder") + self.components.label(top_frame, 0, 4, "placeholder:", + tooltip="The placeholder used when using the embedding in a prompt") + self.components.entry(top_frame, 0, 5, self.ui_state, "placeholder") # token count - components.label(top_frame, 0, 6, "token count:", - tooltip="The token count used when creating a new embedding. Leave empty to auto detect from the initial embedding text.") - token_count_entry = components.entry(top_frame, 0, 7, self.ui_state, "token_count") - token_count_entry.configure(width=40) + self.components.label(top_frame, 0, 6, "token count:", + tooltip="The token count used when creating a new embedding. Leave empty to auto detect from the initial embedding text.") + self.components.entry(top_frame, 0, 7, self.ui_state, "token_count", width=40) # trainable - components.label(bottom_frame, 0, 0, "train:") - trainable_switch = components.switch(bottom_frame, 0, 1, self.ui_state, "train", command=save_command) - trainable_switch.configure(width=40) + self.components.label(bottom_frame, 0, 0, "train:") + self.components.switch(bottom_frame, 0, 1, self.ui_state, "train", command=save_command, width=40) # output embedding - components.label(bottom_frame, 0, 2, "output embedding:", - tooltip="Output embeddings are calculated at the output of the text encoder, not the input. This can improve results for larger text encoders and lower VRAM usage.") - output_embedding_switch = components.switch(bottom_frame, 0, 3, self.ui_state, "is_output_embedding") - output_embedding_switch.configure(width=40) + self.components.label(bottom_frame, 0, 2, "output embedding:", + tooltip="Output embeddings are calculated at the output of the text encoder, not the input. This can improve results for larger text encoders and lower VRAM usage.") + self.components.switch(bottom_frame, 0, 3, self.ui_state, "is_output_embedding", width=40) # stop training after - components.label(bottom_frame, 0, 4, "stop training after:", - tooltip="When to stop training the embedding") - components.time_entry(bottom_frame, 0, 5, self.ui_state, "stop_training_after", "stop_training_after_unit") + self.components.label(bottom_frame, 0, 4, "stop training after:", + tooltip="When to stop training the embedding") + self.components.time_entry(bottom_frame, 0, 5, self.ui_state, "stop_training_after", "stop_training_after_unit") # initial embedding text - components.label(bottom_frame, 0, 6, "initial embedding text:", - tooltip="The initial embedding text used when creating a new embedding") - components.entry(bottom_frame, 0, 7, self.ui_state, "initial_embedding_text") - - def __randomize_uuid(self, embedding_config: TrainEmbeddingConfig): - embedding_config.uuid = TrainEmbeddingConfig.default_values().uuid - return embedding_config + self.components.label(bottom_frame, 0, 6, "initial embedding text:", + tooltip="The initial embedding text used when creating a new embedding") + self.components.entry(bottom_frame, 0, 7, self.ui_state, "initial_embedding_text") def configure_element(self): pass - - def place_in_list(self): - self.grid(row=self.i, column=0, pady=5, padx=5, sticky="new") diff --git a/modules/ui/BaseCaptionUIView.py b/modules/ui/BaseCaptionUIView.py index e6cc0551e..a6561da69 100644 --- a/modules/ui/BaseCaptionUIView.py +++ b/modules/ui/BaseCaptionUIView.py @@ -1,572 +1,32 @@ -import os import platform -import subprocess -import traceback -from tkinter import filedialog -from modules.module.Blip2Model import Blip2Model -from modules.module.BlipModel import BlipModel -from modules.module.ClipSegModel import ClipSegModel -from modules.module.MaskByColor import MaskByColor -from modules.module.RembgHumanModel import RembgHumanModel -from modules.module.RembgModel import RembgModel -from modules.module.WDModel import WDModel -from modules.ui.GenerateCaptionsWindow import GenerateCaptionsWindow -from modules.ui.GenerateMasksWindow import GenerateMasksWindow -from modules.util import path_util -from modules.util.image_util import load_image -from modules.util.torch_util import default_device, torch_gc -from modules.util.ui import components -from modules.util.ui.ui_utils import bind_mousewheel, set_window_icon -from modules.util.ui.UIState import UIState -import torch +class BaseCaptionUIView: + def __init__(self, components): + self.components = components -import customtkinter as ctk -import cv2 -import numpy as np -from customtkinter import ScalingTracker, ThemeManager -from PIL import Image, ImageDraw - - -class CaptionUI(ctk.CTkToplevel): - def __init__( - self, - parent, - initial_dir: str | None, - initial_include_subdirectories: bool, - *args, - **kwargs, - ) -> None: - super().__init__(parent, *args, **kwargs) - self.protocol("WM_DELETE_WINDOW", self._on_close) - - self.dir = initial_dir - self.config_ui_data = {"include_subdirectories": initial_include_subdirectories} - self.config_ui_state = UIState(self, self.config_ui_data) - self.image_size = 850 - self.help_text = """ - Keyboard shortcuts when focusing on the prompt input field: - Up arrow: previous image - Down arrow: next image - Return: save - Ctrl+M: only show the mask - Ctrl+D: draw mask editing mode - Ctrl+F: fill mask editing mode - - When editing masks: - Left click: add mask - Right click: remove mask - Mouse wheel: increase or decrease brush size""" - self.masking_model = None - self.captioning_model = None - self.image_rel_paths = [] - self.current_image_index = -1 - self.file_list = None - self.image_labels = [] - self.pil_image = None - self.image_width = 0 - self.image_height = 0 - self.pil_mask = None - self.mask_draw_x = 0 - self.mask_draw_y = 0 - self.mask_draw_radius = 0.01 - self.display_only_mask = False - self.image = None - self.image_label = None - self.mask_editing_mode = 'draw' - self.enable_mask_editing_var = ctk.BooleanVar() - self.mask_editing_alpha = None - self.prompt_var = None - self.prompt_component = None - - - self.title("OneTrainer") - self.geometry("1280x980") - self.resizable(False, False) - - - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=1) - self.grid_columnconfigure(0, weight=1) - - - self.top_bar(self) - - self.bottom_frame = ctk.CTkFrame(self) - self.bottom_frame.grid(row=1, column=0, sticky="nsew") - self.bottom_frame.grid_rowconfigure(0, weight=1) - self.bottom_frame.grid_columnconfigure(0, weight=0) - self.bottom_frame.grid_columnconfigure(1, weight=1) - - self.file_list_column(self.bottom_frame) - self.content_column(self.bottom_frame) - self.load_directory() - - self.wait_visibility() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - def top_bar(self, master): - top_frame = ctk.CTkFrame(master) - top_frame.grid(row=0, column=0, sticky="nsew") - - components.button(top_frame, 0, 0, "Open", self.open_directory, + def build_top_bar(self, frame, controller, ui_state): + self.components.button(frame, 0, 0, "Open", self.open_directory, tooltip="open a new directory") - components.button(top_frame, 0, 1, "Generate Masks", self.open_mask_window, + self.components.button(frame, 0, 1, "Generate Masks", self.open_mask_window, tooltip="open a dialog to automatically generate masks") - components.button(top_frame, 0, 2, "Generate Captions", self.open_caption_window, + self.components.button(frame, 0, 2, "Generate Captions", self.open_caption_window, tooltip="open a dialog to automatically generate captions") if platform.system() == "Windows": - components.button(top_frame, 0, 3, "Open in Explorer", self.open_in_explorer, + self.components.button(frame, 0, 3, "Open in Explorer", self.open_in_explorer, tooltip="open the current image in Explorer") - components.switch(top_frame, 0, 4, self.config_ui_state, "include_subdirectories", + self.components.switch(frame, 0, 4, ui_state, "include_subdirectories", text="include subdirectories") - top_frame.grid_columnconfigure(5, weight=1) - - components.button(top_frame, 0, 6, "Help", self.print_help, - tooltip=self.help_text) - - def file_list_column(self, master): - if self.file_list is not None: - self.image_labels = [] - self.file_list.destroy() - - self.file_list = ctk.CTkScrollableFrame(master, width=300) - self.file_list.grid(row=0, column=0, sticky="nsew") - - for i, filename in enumerate(self.image_rel_paths): - def __create_switch_image(index): - def __switch_image(event): - self.switch_image(index) - - return __switch_image + frame.grid_columnconfigure(5, weight=1) - label = ctk.CTkLabel(self.file_list, text=filename) - label.bind("", __create_switch_image(i)) + self.components.button(frame, 0, 6, "Help", controller.print_help, + tooltip=controller.help_text) - self.image_labels.append(label) - label.grid(row=i, column=0, padx=5, sticky="nsw") - - def content_column(self, master): - image = Image.new("RGBA", (512, 512), (0, 0, 0, 0)) - - right_frame = ctk.CTkFrame(master, fg_color="transparent") - right_frame.grid(row=0, column=1, sticky="nsew") - - right_frame.grid_columnconfigure(4, weight=1) - right_frame.grid_rowconfigure(1, weight=1) - - components.button(right_frame, 0, 0, "Draw", self.draw_mask_editing_mode, + def build_mask_buttons(self, right_frame): + self.components.button(right_frame, 0, 0, "Draw", self.draw_mask_editing_mode, tooltip="draw a mask using a brush") - components.button(right_frame, 0, 1, "Fill", self.fill_mask_editing_mode, + self.components.button(right_frame, 0, 1, "Fill", self.fill_mask_editing_mode, tooltip="draw a mask using a fill tool") - - # checkbox to enable mask editing - self.enable_mask_editing_var = ctk.BooleanVar() - self.enable_mask_editing_var.set(False) - enable_mask_editing_checkbox = ctk.CTkCheckBox( - right_frame, text="Enable Mask Editing", variable=self.enable_mask_editing_var, width=50) - enable_mask_editing_checkbox.grid(row=0, column=2, padx=25, pady=5, sticky="w") - - # mask alpha textbox - self.mask_editing_alpha = ctk.CTkEntry(master=right_frame, width=40, placeholder_text="1.0") - self.mask_editing_alpha.insert(0, "1.0") - self.mask_editing_alpha.grid(row=0, column=3, sticky="e", padx=5, pady=5) - self.bind_key_events(self.mask_editing_alpha) - - mask_editing_alpha_label = ctk.CTkLabel(right_frame, text="Brush Alpha", width=75) - mask_editing_alpha_label.grid(row=0, column=4, padx=0, pady=5, sticky="w") - - # image - self.image = ctk.CTkImage( - light_image=image, - size=(self.image_size, self.image_size) - ) - self.image_label = ctk.CTkLabel( - master=right_frame, text="", image=self.image, height=self.image_size, width=self.image_size - ) - self.image_label.grid(row=1, column=0, columnspan=5, sticky="nsew") - - self.image_label.bind("", self.edit_mask) - self.image_label.bind("", self.edit_mask) - self.image_label.bind("", self.edit_mask) - bind_mousewheel(self.image_label, {self.image_label.children["!label"]}, self.draw_mask_radius) - - # prompt - self.prompt_var = ctk.StringVar() - self.prompt_component = ctk.CTkEntry(right_frame, textvariable=self.prompt_var) - self.prompt_component.grid(row=2, column=0, columnspan=5, pady=5, sticky="new") - self.bind_key_events(self.prompt_component) - self.prompt_component.focus_set() - - def bind_key_events(self, component): - component.bind("", self.next_image) - component.bind("", self.previous_image) - component.bind("", self.save) - component.bind("", self.toggle_mask) - component.bind("", self.draw_mask_editing_mode) - component.bind("", self.fill_mask_editing_mode) - - def load_directory(self, include_subdirectories: bool = False): - self.scan_directory(include_subdirectories) - self.file_list_column(self.bottom_frame) - - if len(self.image_rel_paths) > 0: - self.switch_image(0) - else: - self.switch_image(-1) - - self.prompt_component.focus_set() - - def scan_directory(self, include_subdirectories: bool = False): - def __is_supported_image_extension(filename): - name, ext = os.path.splitext(filename) - return path_util.is_supported_image_extension(ext) and not name.endswith("-masklabel") and not name.endswith("-condlabel") - - self.image_rel_paths = [] - - if not self.dir or not os.path.isdir(self.dir): - return - - if include_subdirectories: - for root, _, files in os.walk(self.dir): - for filename in files: - if __is_supported_image_extension(filename): - self.image_rel_paths.append( - os.path.relpath(os.path.join(root, filename), self.dir) - ) - else: - for _, filename in enumerate(os.listdir(self.dir)): - if __is_supported_image_extension(filename): - self.image_rel_paths.append( - os.path.relpath(os.path.join(self.dir, filename), self.dir) - ) - - def load_image(self): - image_name = "resources/icons/icon.png" - - if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): - image_name = self.image_rel_paths[self.current_image_index] - image_name = os.path.join(self.dir, image_name) - - try: - return load_image(image_name, convert_mode="RGB") - except Exception: - print(f'Could not open image {image_name}') - - def load_mask(self): - if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): - image_name = self.image_rel_paths[self.current_image_index] - mask_name = os.path.splitext(image_name)[0] + "-masklabel.png" - mask_name = os.path.join(self.dir, mask_name) - - try: - return load_image(mask_name, convert_mode='RGB') - except Exception: - return None - else: - return None - - def load_prompt(self): - if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): - image_name = self.image_rel_paths[self.current_image_index] - prompt_name = os.path.splitext(image_name)[0] + ".txt" - prompt_name = os.path.join(self.dir, prompt_name) - - try: - with open(prompt_name, "r", encoding='utf-8') as f: - return f.readlines()[0].strip() - except Exception: - return "" - else: - return "" - - def previous_image(self, event): - if len(self.image_rel_paths) > 0 and (self.current_image_index - 1) >= 0: - self.switch_image(self.current_image_index - 1) - - def next_image(self, event): - if len(self.image_rel_paths) > 0 and (self.current_image_index + 1) < len(self.image_rel_paths): - self.switch_image(self.current_image_index + 1) - - def switch_image(self, index): - if len(self.image_labels) > 0 and self.current_image_index < len(self.image_labels): - self.image_labels[self.current_image_index].configure( - text_color=ThemeManager.theme["CTkLabel"]["text_color"]) - - self.current_image_index = index - if index >= 0: - self.image_labels[index].configure(text_color="#FF0000") - - self.pil_image = self.load_image() - self.pil_mask = self.load_mask() - prompt = self.load_prompt() - - self.image_width = self.pil_image.width - self.image_height = self.pil_image.height - scale = self.image_size / max(self.pil_image.height, self.pil_image.width) - height = int(self.pil_image.height * scale) - width = int(self.pil_image.width * scale) - - self.pil_image = self.pil_image.resize((width, height), Image.Resampling.LANCZOS) - - self.refresh_image() - self.prompt_var.set(prompt) - else: - image = Image.new("RGB", (512, 512), (0, 0, 0)) - self.image.configure(light_image=image) - - def refresh_image(self): - if self.pil_mask: - resized_pil_mask = self.pil_mask.resize( - (self.pil_image.width, self.pil_image.height), - Image.Resampling.NEAREST - ) - - if self.display_only_mask: - self.image.configure(light_image=resized_pil_mask, size=resized_pil_mask.size) - else: - np_image = np.array(self.pil_image).astype(np.float32) / 255.0 - np_mask = np.array(resized_pil_mask).astype(np.float32) / 255.0 - - # normalize mask between 0.3 - 1.0 so we can see image underneath and gauge strength of the alpha - norm_min = 0.3 - np_mask_min = np_mask.min() - if np_mask_min == 0: - # optimize for common case - np_mask = np_mask * (1.0 - norm_min) + norm_min - elif np_mask_min < 1: - # note: min of 1 means we get divide by 0 - np_mask = (np_mask - np_mask_min) / (1.0 - np_mask_min) * (1.0 - norm_min) + norm_min - - np_masked_image = (np_image * np_mask * 255.0).astype(np.uint8) - masked_image = Image.fromarray(np_masked_image, mode='RGB') - - self.image.configure(light_image=masked_image, size=masked_image.size) - else: - self.image.configure(light_image=self.pil_image, size=self.pil_image.size) - - def draw_mask_radius(self, delta, raw_event): - # Wheel up = Increase radius. Wheel down = Decrease radius. - multiplier = 1.0 + (delta * 0.05) - self.mask_draw_radius = max(0.0025, self.mask_draw_radius * multiplier) - - def edit_mask(self, event): - if not self.enable_mask_editing_var.get(): - return - - if event.widget != self.image_label.children["!label"]: - return - - if len(self.image_rel_paths) == 0 or self.current_image_index >= len(self.image_rel_paths): - return - - display_scaling = ScalingTracker.get_window_scaling(self) - - event_x = event.x / display_scaling - event_y = event.y / display_scaling - - start_x = int(event_x / self.pil_image.width * self.image_width) - start_y = int(event_y / self.pil_image.height * self.image_height) - end_x = int(self.mask_draw_x / self.pil_image.width * self.image_width) - end_y = int(self.mask_draw_y / self.pil_image.height * self.image_height) - - self.mask_draw_x = event_x - self.mask_draw_y = event_y - - is_right = False - is_left = False - if event.state & 0x0100 or event.num == 1: # left mouse button - is_left = True - elif event.state & 0x0400 or event.num == 3: # right mouse button - is_right = True - - if self.mask_editing_mode == 'draw': - self.draw_mask(start_x, start_y, end_x, end_y, is_left, is_right) - if self.mask_editing_mode == 'fill': - self.fill_mask(start_x, start_y, end_x, end_y, is_left, is_right) - - def draw_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): - color = None - - adding_to_mask = True - if is_left: - try: - alpha = float(self.mask_editing_alpha.get()) - except Exception: - alpha = 1.0 - rgb_value = int(max(0.0, min(alpha, 1.0)) * 255) # max/min stuff to clamp to 0 - 255 range - color = (rgb_value, rgb_value, rgb_value) - - elif is_right: - color = (0, 0, 0) - adding_to_mask = False - - if color is not None: - if self.pil_mask is None: - if adding_to_mask: - self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(0, 0, 0)) - else: - self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(255, 255, 255)) - - radius = int(self.mask_draw_radius * max(self.pil_mask.width, self.pil_mask.height)) - - draw = ImageDraw.Draw(self.pil_mask) - draw.line((start_x, start_y, end_x, end_y), fill=color, - width=radius + radius + 1) - draw.ellipse((start_x - radius, start_y - radius, - start_x + radius, start_y + radius), fill=color, outline=None) - draw.ellipse((end_x - radius, end_y - radius, end_x + radius, - end_y + radius), fill=color, outline=None) - - self.refresh_image() - - def fill_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): - color = None - - adding_to_mask = True - if is_left: - try: - alpha = float(self.mask_editing_alpha.get()) - except Exception: - alpha = 1.0 - rgb_value = int(max(0.0, min(alpha, 1.0)) * 255) # max/min stuff to clamp to 0 - 255 range - color = (rgb_value, rgb_value, rgb_value) - - elif is_right: - color = (0, 0, 0) - adding_to_mask = False - - if color is not None: - if self.pil_mask is None: - if adding_to_mask: - self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(0, 0, 0)) - else: - self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(255, 255, 255)) - - np_mask = np.array(self.pil_mask).astype(np.uint8) - cv2.floodFill(np_mask, None, (start_x, start_y), color) - self.pil_mask = Image.fromarray(np_mask, 'RGB') - - self.refresh_image() - - def save(self, event): - if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): - image_name = self.image_rel_paths[self.current_image_index] - - prompt_name = os.path.splitext(image_name)[0] + ".txt" - prompt_name = os.path.join(self.dir, prompt_name) - - mask_name = os.path.splitext(image_name)[0] + "-masklabel.png" - mask_name = os.path.join(self.dir, mask_name) - - try: - with open(prompt_name, "w", encoding='utf-8') as f: - f.write(self.prompt_var.get()) - except Exception: - return - - if self.pil_mask: - self.pil_mask.save(mask_name) - - def draw_mask_editing_mode(self, *args): - self.mask_editing_mode = 'draw' - - if args: - # disable default event - return "break" - return None - - def fill_mask_editing_mode(self, *args): - self.mask_editing_mode = 'fill' - - def toggle_mask(self, *args): - self.display_only_mask = not self.display_only_mask - self.refresh_image() - - def open_directory(self): - new_dir = filedialog.askdirectory() - - if new_dir: - self.dir = new_dir - self.load_directory(include_subdirectories=self.config_ui_data["include_subdirectories"]) - - def open_mask_window(self): - dialog = GenerateMasksWindow(self, self.dir, self.config_ui_data["include_subdirectories"]) - self.wait_window(dialog) - self.switch_image(self.current_image_index) - - def open_caption_window(self): - dialog = GenerateCaptionsWindow(self, self.dir, self.config_ui_data["include_subdirectories"]) - self.wait_window(dialog) - self.switch_image(self.current_image_index) - - def open_in_explorer(self): - try: - image_name = self.image_rel_paths[self.current_image_index] - image_name = os.path.realpath(os.path.join(self.dir, image_name)) - subprocess.Popen(f"explorer /select,{image_name}") - except Exception: - traceback.print_exc() - - def load_masking_model(self, model): - model_type = type(self.masking_model).__name__ if self.masking_model else None - - if model == "ClipSeg" and model_type != "ClipSegModel": - self._release_models() - print("loading ClipSeg model, this may take a while") - self.masking_model = ClipSegModel(default_device, torch.float32) - elif model == "Rembg" and model_type != "RembgModel": - self._release_models() - print("loading Rembg model, this may take a while") - self.masking_model = RembgModel(default_device, torch.float32) - elif model == "Rembg-Human" and model_type != "RembgHumanModel": - self._release_models() - print("loading Rembg-Human model, this may take a while") - self.masking_model = RembgHumanModel(default_device, torch.float32) - elif model == "Hex Color" and model_type != "MaskByColor": - self._release_models() - self.masking_model = MaskByColor(default_device, torch.float32) - - def load_captioning_model(self, model): - model_type = type(self.captioning_model).__name__ if self.captioning_model else None - - if model == "Blip" and model_type != "BlipModel": - self._release_models() - print("loading Blip model, this may take a while") - self.captioning_model = BlipModel(default_device, torch.float16) - elif model == "Blip2" and model_type != "Blip2Model": - self._release_models() - print("loading Blip2 model, this may take a while") - self.captioning_model = Blip2Model(default_device, torch.float16) - elif model == "WD14 VIT v2" and model_type != "WDModel": - self._release_models() - print("loading WD14_VIT_v2 model, this may take a while") - self.captioning_model = WDModel(default_device, torch.float16) - - def print_help(self): - print(self.help_text) - - def _release_models(self): - """Release all models from VRAM""" - freed = False - if self.captioning_model is not None: - self.captioning_model = None - freed = True - if self.masking_model is not None: - self.masking_model = None - freed = True - if freed: - torch_gc() - - def _on_close(self): - self._release_models() - self.destroy() - - def destroy(self): - self._release_models() - super().destroy() diff --git a/modules/ui/BaseCloudTabView.py b/modules/ui/BaseCloudTabView.py index 99057e428..80be17360 100644 --- a/modules/ui/BaseCloudTabView.py +++ b/modules/ui/BaseCloudTabView.py @@ -1,221 +1,184 @@ -import webbrowser +from abc import ABC, abstractmethod -from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.CloudAction import CloudAction from modules.util.enum.CloudFileSync import CloudFileSync from modules.util.enum.CloudType import CloudType -from modules.util.ui import components -from modules.util.ui.UIState import UIState -import customtkinter as ctk +class BaseCloudTabView(ABC): + def __init__(self, components): + self.components = components -class CloudTab: + @property + def reattach(self): + return self.controller.reattach - def __init__(self, master, train_config: TrainConfig, ui_state: UIState, parent): - super().__init__() + @abstractmethod + def _make_reattach_frame(self, frame): pass - self.master = master - self.train_config = train_config - self.ui_state = ui_state - self.parent = parent - self.reattach = False + @abstractmethod + def _make_create_frame(self, frame): pass - self.frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) - self.frame.grid_columnconfigure(2, weight=0) - self.frame.grid_columnconfigure(3, weight=1) - self.frame.grid_columnconfigure(4, weight=0) - self.frame.grid_columnconfigure(5, weight=1) + @abstractmethod + def _on_set_gpu_types(self): pass - components.label(self.frame, 0, 0, "Enabled", + def build_content(self, frame, controller, ui_state): + self.components.label(frame, 0, 0, "Enabled", tooltip="Enable cloud training") - components.switch(self.frame, 0, 1, self.ui_state, "cloud.enabled") + self.components.switch(frame, 0, 1, ui_state, "cloud.enabled") - components.label(self.frame, 1, 0, "Type", + self.components.label(frame, 1, 0, "Type", tooltip="Choose LINUX to connect to a linux machine via SSH. Choose RUNPOD for additional functionality such as automatically creating and deleting pods.") - components.options_kv(self.frame, 1, 1, [ + self.components.options_kv(frame, 1, 1, [ ("RUNPOD", CloudType.RUNPOD), ("LINUX", CloudType.LINUX), - ], self.ui_state, "cloud.type") + ], ui_state, "cloud.type") - components.label(self.frame, 2, 0, "File sync method", + self.components.label(frame, 2, 0, "File sync method", tooltip="Choose NATIVE_SCP to use scp.exe to transfer files. FABRIC_SFTP uses the Paramiko/Fabric SFTP implementation for file transfers instead.") - components.options_kv(self.frame, 2, 1, [ + self.components.options_kv(frame, 2, 1, [ ("NATIVE_SCP", CloudFileSync.NATIVE_SCP), ("FABRIC_SFTP", CloudFileSync.FABRIC_SFTP), - ], self.ui_state, "cloud.file_sync") + ], ui_state, "cloud.file_sync") - components.label(self.frame, 3, 0, "API key", + self.components.label(frame, 3, 0, "API key", tooltip="Cloud service API key for RUNPOD. Leave empty for LINUX. This value is stored separately, not saved to your configuration file. ") - components.entry(self.frame, 3, 1, self.ui_state, "secrets.cloud.api_key") + self.components.entry(frame, 3, 1, ui_state, "secrets.cloud.api_key") - components.label(self.frame, 4, 0, "Hostname", + self.components.label(frame, 4, 0, "Hostname", tooltip="SSH server hostname or IP. Leave empty if you have a Cloud ID or want to automatically create a new cloud.") - components.entry(self.frame, 4, 1, self.ui_state, "secrets.cloud.host") + self.components.entry(frame, 4, 1, ui_state, "secrets.cloud.host") - components.label(self.frame, 5, 0, "Port", + self.components.label(frame, 5, 0, "Port", tooltip="SSH server port. Leave empty if you have a Cloud ID or want to automatically create a new cloud.") - components.entry(self.frame, 5, 1, self.ui_state, "secrets.cloud.port") + self.components.entry(frame, 5, 1, ui_state, "secrets.cloud.port") - components.label(self.frame, 6, 0, "User", + self.components.label(frame, 6, 0, "User", tooltip='SSH username. Use "root" for RUNPOD. Your SSH client must be set up to connect to the cloud using a public key, without a password. For RUNPOD, create an ed25519 key locally, and copy the contents of the public keyfile to your "SSH Public Keys" on the RunPod website.') - components.entry(self.frame, 6, 1, self.ui_state, "secrets.cloud.user") + self.components.entry(frame, 6, 1, ui_state, "secrets.cloud.user") - components.label(self.frame, 7, 0, "SSH keyfile path", + self.components.label(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.path_entry(self.frame, 7, 1, self.ui_state, "secrets.cloud.key_file", mode="file") + self.components.path_entry(frame, 7, 1, ui_state, "secrets.cloud.key_file", mode="file") - components.label(self.frame, 8, 0, "SSH password", + self.components.label(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.") - components.entry(self.frame, 8, 1, self.ui_state, "secrets.cloud.password") + self.components.entry(frame, 8, 1, ui_state, "secrets.cloud.password") - components.label(self.frame, 9, 0, "Cloud id", + self.components.label(frame, 9, 0, "Cloud id", tooltip="RUNPOD Cloud ID. The cloud service must have a public IP and SSH service. Leave empty if you want to automatically create a new RUNPOD cloud, or if you're connecting to another cloud provider via SSH Hostname and Port.") - components.entry(self.frame, 9, 1, self.ui_state, "secrets.cloud.id") + self.components.entry(frame, 9, 1, ui_state, "secrets.cloud.id") - components.label(self.frame, 10, 0, "Tensorboard TCP tunnel", + self.components.label(frame, 10, 0, "Tensorboard TCP tunnel", tooltip="Instead of starting tensorboard locally, make a TCP tunnel to a tensorboard on the cloud") - components.switch(self.frame, 10, 1, self.ui_state, "cloud.tensorboard_tunnel") + self.components.switch(frame, 10, 1, ui_state, "cloud.tensorboard_tunnel") - - - components.label(self.frame, 1, 2, "Remote Directory", + self.components.label(frame, 1, 2, "Remote Directory", tooltip="The directory on the cloud where files will be uploaded and downloaded.") - components.entry(self.frame, 1, 3, self.ui_state, "cloud.remote_dir") - components.label(self.frame, 2, 2, "OneTrainer Directory", + self.components.entry(frame, 1, 3, ui_state, "cloud.remote_dir") + self.components.label(frame, 2, 2, "OneTrainer Directory", tooltip="The directory for OneTrainer on the cloud.") - components.entry(self.frame, 2, 3, self.ui_state, "cloud.onetrainer_dir") - components.label(self.frame, 3, 2, "Huggingface cache Directory", + self.components.entry(frame, 2, 3, ui_state, "cloud.onetrainer_dir") + self.components.label(frame, 3, 2, "Huggingface cache Directory", tooltip="Huggingface models are downloaded to this remote directory.") - components.entry(self.frame, 3, 3, self.ui_state, "cloud.huggingface_cache_dir") - components.label(self.frame, 4, 2, "Install OneTrainer", + self.components.entry(frame, 3, 3, ui_state, "cloud.huggingface_cache_dir") + self.components.label(frame, 4, 2, "Install OneTrainer", tooltip="Automatically install OneTrainer from GitHub if the directory doesn't already exist.") - components.switch(self.frame, 4, 3, self.ui_state, "cloud.install_onetrainer") - components.label(self.frame, 5, 2, "Install command", + self.components.switch(frame, 4, 3, ui_state, "cloud.install_onetrainer") + self.components.label(frame, 5, 2, "Install command", tooltip="The command for installing OneTrainer. Leave the default, unless you want to use a development branch of OneTrainer.") - components.entry(self.frame, 5, 3, self.ui_state, "cloud.install_cmd") - components.label(self.frame, 6, 2, "Update OneTrainer", + self.components.entry(frame, 5, 3, ui_state, "cloud.install_cmd") + self.components.label(frame, 6, 2, "Update OneTrainer", tooltip="Update OneTrainer if it already exists on the cloud.") - components.switch(self.frame, 6, 3, self.ui_state, "cloud.update_onetrainer") + self.components.switch(frame, 6, 3, ui_state, "cloud.update_onetrainer") - components.label(self.frame, 8, 2, "Detach remote trainer", + self.components.label(frame, 8, 2, "Detach remote trainer", tooltip="Allows the trainer to keep running even if your connection to the cloud is lost.") - components.switch(self.frame, 8, 3, self.ui_state, "cloud.detach_trainer") - components.label(self.frame, 9, 2, "Reattach id", + self.components.switch(frame, 8, 3, ui_state, "cloud.detach_trainer") + self.components.label(frame, 9, 2, "Reattach id", tooltip="An id identifying the remotely running trainer. In case you have lost connection or closed OneTrainer, it will try to reattach to this id instead of starting a new remote trainer.") - reattach_frame = ctk.CTkFrame(self.frame, fg_color="transparent") - reattach_frame.grid(row=9, column=3, padx=0, pady=0, sticky="new") - reattach_frame.grid_columnconfigure(0, weight=1) - reattach_frame.grid_columnconfigure(1, weight=1) - components.entry(reattach_frame, 0, 0, self.ui_state, "cloud.run_id", width=60) - components.button(reattach_frame, 0, 1, "Reattach now", self.__reattach) - - components.label(self.frame, 11, 2, "Download samples", + reattach_frame = self._make_reattach_frame(frame) + self.components.entry(reattach_frame, 0, 0, ui_state, "cloud.run_id", width=60) + self.components.button(reattach_frame, 0, 1, "Reattach now", controller.do_reattach) + + self.components.label(frame, 11, 2, "Download samples", tooltip="Download samples from the remote workspace directory to your local machine.") - components.switch(self.frame, 11, 3, self.ui_state, "cloud.download_samples") - components.label(self.frame, 12, 2, "Download output model", + self.components.switch(frame, 11, 3, ui_state, "cloud.download_samples") + self.components.label(frame, 12, 2, "Download output model", tooltip="Download the final model after training. You can disable this if you plan to use an automatically saved checkpoint instead.") - components.switch(self.frame, 12, 3, self.ui_state, "cloud.download_output_model") - components.label(self.frame, 13, 2, "Download saved checkpoints", + self.components.switch(frame, 12, 3, ui_state, "cloud.download_output_model") + self.components.label(frame, 13, 2, "Download saved checkpoints", tooltip="Download the automatically saved training checkpoints from the remote workspace directory to your local machine.") - components.switch(self.frame, 13, 3, self.ui_state, "cloud.download_saves") - components.label(self.frame, 14, 2, "Download backups", + self.components.switch(frame, 13, 3, ui_state, "cloud.download_saves") + self.components.label(frame, 14, 2, "Download backups", tooltip="Download backups from the remote workspace directory to your local machine. It's usually not necessary to download them, because as long as the backups are still available on the cloud, the training can be restarted using one of the cloud's backups.") - components.switch(self.frame, 14, 3, self.ui_state, "cloud.download_backups") - components.label(self.frame, 15, 2, "Download tensorboard logs", + self.components.switch(frame, 14, 3, ui_state, "cloud.download_backups") + self.components.label(frame, 15, 2, "Download tensorboard logs", tooltip="Download TensorBoard event logs from the remote workspace directory to your local machine. They can then be viewed locally in TensorBoard. It is recommended to disable \"Sample to TensorBoard\" to reduce the event log size.") - components.switch(self.frame, 15, 3, self.ui_state, "cloud.download_tensorboard") - components.label(self.frame, 16, 2, "Delete remote workspace", + self.components.switch(frame, 15, 3, ui_state, "cloud.download_tensorboard") + self.components.label(frame, 16, 2, "Delete remote workspace", tooltip="Delete the workspace directory on the cloud after training has finished successfully and data has been downloaded.") - components.switch(self.frame, 16, 3, self.ui_state, "cloud.delete_workspace") + self.components.switch(frame, 16, 3, ui_state, "cloud.delete_workspace") - components.label(self.frame, 1, 4, "Create cloud via API", + self.components.label(frame, 1, 4, "Create cloud via API", tooltip="Automatically creates a new cloud instance if both Host:Port and Cloud ID are empty. Currently supported for RUNPOD.") - create_frame = ctk.CTkFrame(self.frame, fg_color="transparent") - create_frame.grid(row=1, column=5, padx=0, pady=0, sticky="new") - create_frame.grid_columnconfigure(0, weight=0) - create_frame.grid_columnconfigure(1, weight=1) - components.switch(create_frame, 0, 0, self.ui_state, "cloud.create") - components.button(create_frame, 0, 1, "Create cloud via website", self.__create_cloud) - - components.label(self.frame, 2, 4, "Cloud name", + create_frame = self._make_create_frame(frame) + self.components.switch(create_frame, 0, 0, ui_state, "cloud.create") + self.components.button(create_frame, 0, 1, "Create cloud via website", controller.open_create_cloud_url) + + self.components.label(frame, 2, 4, "Cloud name", tooltip="The name of the new cloud instance.") - components.entry(self.frame, 2, 5, self.ui_state, "cloud.name") - components.label(self.frame, 3, 4, "Type", + self.components.entry(frame, 2, 5, ui_state, "cloud.name") + self.components.label(frame, 3, 4, "Type", tooltip="Select the RunPod cloud type. See RunPod's website for details.") - components.options_kv(self.frame, 3, 5, [ + self.components.options_kv(frame, 3, 5, [ ("", ""), ("Community", "COMMUNITY"), ("Secure", "SECURE"), - ], self.ui_state, "cloud.sub_type") - + ], ui_state, "cloud.sub_type") - components.label(self.frame, 4, 4, "GPU", + self.components.label(frame, 4, 4, "GPU", tooltip="Select the GPU type. Enter an API key before pressing the button.") + _, gpu_components = self.components.options_adv(frame, 4, 5, [("")], ui_state, "cloud.gpu_type", adv_command=self._on_set_gpu_types) + self.gpu_types_menu = gpu_components['component'] - _,gpu_components=components.options_adv(self.frame, 4, 5, [("")], self.ui_state, "cloud.gpu_type",adv_command=self.__set_gpu_types) - self.gpu_types_menu=gpu_components['component'] - - components.label(self.frame, 5, 4, "Volume size", + self.components.label(frame, 5, 4, "Volume size", tooltip="Set the storage volume size in GB. This volume persists only until the cloud is deleted - not a RunPod network volume") - components.entry(self.frame, 5, 5, self.ui_state, "cloud.volume_size") + self.components.entry(frame, 5, 5, ui_state, "cloud.volume_size") - components.label(self.frame, 6, 4, "Min download", + self.components.label(frame, 6, 4, "Min download", tooltip="Set the minimum download speed of the cloud in Mbps.") - components.entry(self.frame, 6, 5, self.ui_state, "cloud.min_download") + self.components.entry(frame, 6, 5, ui_state, "cloud.min_download") - components.label(self.frame, 8, 4, "Action on finish", + self.components.label(frame, 8, 4, "Action on finish", tooltip="What to do when training finishes and the data has been fully downloaded: Stop or delete the cloud, or do nothing.") - components.options_kv(self.frame, 8, 5, [ + self.components.options_kv(frame, 8, 5, [ ("None", CloudAction.NONE), ("Stop", CloudAction.STOP), ("Delete", CloudAction.DELETE), - ], self.ui_state, "cloud.on_finish") + ], ui_state, "cloud.on_finish") - components.label(self.frame, 9, 4, "Action on error", + self.components.label(frame, 9, 4, "Action on error", tooltip="What to do if training stops due to an error: Stop or delete the cloud, or do nothing. Data may be lost.") - components.options_kv(self.frame, 9, 5, [ + self.components.options_kv(frame, 9, 5, [ ("None", CloudAction.NONE), ("Stop", CloudAction.STOP), ("Delete", CloudAction.DELETE), - ], self.ui_state, "cloud.on_error") + ], ui_state, "cloud.on_error") - components.label(self.frame, 10, 4, "Action on detached finish", + self.components.label(frame, 10, 4, "Action on detached finish", tooltip="What to do when training finishes, but the client has been detached and cannot download data. Data may be lost.") - components.options_kv(self.frame, 10, 5, [ + self.components.options_kv(frame, 10, 5, [ ("None", CloudAction.NONE), ("Stop", CloudAction.STOP), ("Delete", CloudAction.DELETE), - ], self.ui_state, "cloud.on_detached_finish") + ], ui_state, "cloud.on_detached_finish") - components.label(self.frame, 11, 4, "Action on detached error", + self.components.label(frame, 11, 4, "Action on detached error", tooltip="What to if training stops due to an error, but the client has been detached and cannot download data. Data may be lost.") - components.options_kv(self.frame, 11, 5, [ + self.components.options_kv(frame, 11, 5, [ ("None", CloudAction.NONE), ("Stop", CloudAction.STOP), ("Delete", CloudAction.DELETE), - ], self.ui_state, "cloud.on_detached_error") - - self.frame.pack(fill="both", expand=1) - - def __set_gpu_types(self): - self.gpu_types_menu.configure(values=[]) - if self.train_config.cloud.type == CloudType.RUNPOD: - import runpod - runpod.api_key=self.train_config.secrets.cloud.api_key - gpus=runpod.get_gpus() - self.gpu_types_menu.configure(values=[gpu['id'] for gpu in gpus]) - - def __reattach(self): - self.reattach=True - try: - self.parent.start_training() - finally: - self.reattach=False - - def __create_cloud(self): - if self.train_config.cloud.type == CloudType.RUNPOD: - webbrowser.open("https://www.runpod.io/console/deploy?template=1a33vbssq9&type=gpu", new=0, autoraise=False) + ], ui_state, "cloud.on_detached_error") diff --git a/modules/ui/BaseConceptTabView.py b/modules/ui/BaseConceptTabView.py index 0b6505694..077a10c24 100644 --- a/modules/ui/BaseConceptTabView.py +++ b/modules/ui/BaseConceptTabView.py @@ -1,112 +1,24 @@ import os import pathlib -from tkinter import BooleanVar, StringVar -from modules.ui.ConceptWindow import ConceptWindow -from modules.ui.ConfigList import ConfigList +from modules.ui.BaseConfigListView import BaseConfigListView +from modules.ui.ConceptWindowController import ConceptWindowController from modules.util import path_util -from modules.util.config.ConceptConfig import ConceptConfig -from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.ConceptType import ConceptType from modules.util.image_util import load_image -from modules.util.ui import components -from modules.util.ui.UIState import UIState -from modules.util.ui.validation import DebounceTimer -import customtkinter as ctk from PIL import Image -class ConceptTab(ConfigList): +class BaseConceptTabView(BaseConfigListView): - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - self.search_var = StringVar() - self.filter_var = StringVar(value="ALL") - self.show_disabled_var = BooleanVar(value=True) - - super().__init__( - master, - train_config, - ui_state, - from_external_file=True, - attr_name="concept_file_name", - config_dir="training_concepts", - default_config_name="concepts.json", - add_button_text="Add Concept", - add_button_tooltip="Adds a new concept to the current config.", - is_full_width=False, - show_toggle_button=True - ) - self._toolbar = None - self._toolbar_is_wrapped = False - self._add_search_bar() - # wrap toolbar if too narrow - self.top_frame.bind('', lambda e: self._maybe_reposition_toolbar(e.width)) - - def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - return ConceptWidget(master, element, i, open_command, remove_command, clone_command, save_command) - - def create_new_element(self) -> dict: - return ConceptConfig.default_values() - - def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: - return ConceptWindow(self.master, self.train_config, self.current_config[i], ui_state[0], ui_state[1], ui_state[2]) - - def _add_search_bar(self): - toolbar = ctk.CTkFrame(self.top_frame, fg_color="transparent") - toolbar.grid(row=0, column=4, columnspan=2, padx=10, sticky="ew") - toolbar.grid_columnconfigure(2, weight=1) - self._toolbar = toolbar - - # Search - ctk.CTkLabel(toolbar, text="Search:").grid(row=0, column=0, padx=(0,5)) - self.search_var = StringVar() - self.search_entry = ctk.CTkEntry(toolbar, textvariable=self.search_var, - placeholder_text="Filter...", width=200) - self.search_entry.grid(row=0, column=1) - self._search_debouncer = DebounceTimer(self.search_entry, 300, lambda: self._update_filters()) - self.search_var.trace_add("write", lambda *_: self._search_debouncer.call()) - - # Spacer - ctk.CTkLabel(toolbar, text="").grid(row=0, column=2, padx=5) - - # Type filter - ctk.CTkLabel(toolbar, text="Type:").grid(row=0, column=3, padx=(0,5)) - self.filter_var = StringVar(value="ALL") - ctk.CTkOptionMenu(toolbar, values=["ALL", "STANDARD", "VALIDATION", "PRIOR_PREDICTION"], - variable=self.filter_var, command=lambda x: self._update_filters(), - width=150).grid(row=0, column=4) - - # Show disabled checkbox - self.show_disabled_var = BooleanVar(value=True) - self.show_disabled_checkbox = ctk.CTkCheckBox(toolbar, text="Show Disabled", variable=self.show_disabled_var, - command=self._update_filters, width=100) - self.show_disabled_checkbox.grid(row=0, column=5, padx=(10,0)) - self._refresh_show_disabled_text() - - # Clear button - ctk.CTkButton(toolbar, text="Clear", width=50, - command=self._reset_filters).grid(row=0, column=6, padx=(10,0)) - - def _update_filters(self): - self._create_element_list(search=self.search_var.get(), - type=self.filter_var.get(), - show_disabled=self.show_disabled_var.get()) - self._refresh_show_disabled_text() - - def _reset_filters(self): - self.search_var.set("") - self.filter_var.set("ALL") - self.show_disabled_var.set(True) - self._update_filters() + _FILTER_TYPES = ["ALL", "STANDARD", "VALIDATION", "PRIOR_PREDICTION"] def _element_matches_filters(self, element): - # Check enabled status if not self.filters.get("show_disabled", True): if hasattr(element, 'enabled') and not element.enabled: return False - # Search filter search = self.filters.get("search", "").lower() if search: if not hasattr(element, '_search_cache'): @@ -127,7 +39,6 @@ def _element_matches_filters(self, element): if not any(search in text for text in getattr(element, '_search_cache', [])): return False - # Type filter type_filter = self.filters.get("type", "ALL") if type_filter != "ALL": if hasattr(element, 'type') and element.type: @@ -139,101 +50,19 @@ def _element_matches_filters(self, element): return True - def _maybe_reposition_toolbar(self, width): - if not self._toolbar: - return - threshold = 1070 - want_wrapped = width < threshold - if want_wrapped == self._toolbar_is_wrapped: - return - self._toolbar_is_wrapped = want_wrapped - if want_wrapped: - self._toolbar.grid_configure(row=1, column=0, columnspan=8, sticky="ew", padx=10) - else: - self._toolbar.grid_configure(row=0, column=4, columnspan=2, sticky="ew", padx=10) - - def _refresh_show_disabled_text(self): - try: - disabled_count = sum(1 for c in getattr(self, 'current_config', []) if getattr(c, 'enabled', True) is False) - except (AttributeError, TypeError): - disabled_count = 0 - text = f"Show Disabled ({disabled_count})" if disabled_count > 0 else "Show Disabled" - try: - if getattr(self, 'show_disabled_checkbox', None): - self.show_disabled_checkbox.configure(text=text) - except (AttributeError, RuntimeError): - pass - - -class ConceptWidget(ctk.CTkFrame): - def __init__(self, master, concept, i, open_command, remove_command, clone_command, save_command): - super().__init__( - master=master, width=150, height=170, corner_radius=10, bg_color="transparent" - ) - - self.concept = concept - self.ui_state = UIState(self, concept) - self.image_ui_state = UIState(self, concept.image) - self.text_ui_state = UIState(self, concept.text) - self.i = i - - self.grid_rowconfigure(1, weight=1) - - # image - self.image = ctk.CTkImage( - light_image=self.__get_preview_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) - - # name - self.name_label = components.label(self, 1, 0, self.__get_display_name(), pad=5, wraplength=140) - - # close button - close_button = ctk.CTkButton( - master=self, - width=20, - height=20, - text="X", - corner_radius=2, - fg_color="#C00000", - command=lambda: remove_command(self.i), - ) - close_button.place(x=0, y=0) - - # clone button - clone_button = ctk.CTkButton( - master=self, - width=20, - height=20, - text="+", - corner_radius=2, - fg_color="#00C000", - command=lambda: clone_command(self.i, self.__randomize_seed), - ) - clone_button.place(x=25, y=0) + def _update_filters(self): + self._create_element_list(search=self.search_var.get(), + type=self.filter_var.get(), + show_disabled=self.show_disabled_var.get()) + self._refresh_show_disabled_text() - # enabled switch - enabled_switch = ctk.CTkSwitch( - master=self, - width=40, - variable=self.ui_state.get_var("enabled"), - text="", - command=save_command, - ) - enabled_switch.place(x=110, y=0) - image_label.bind( - "", - lambda event: open_command(self.i, (self.ui_state, self.image_ui_state, self.text_ui_state)) - ) +class BaseConceptWidgetView: - def __randomize_seed(self, concept: ConceptConfig): - concept.seed = ConceptConfig.default_values().seed - return concept + def __init__(self, components): + self.components = components - def __get_display_name(self): + def _get_display_name(self): if self.concept.name: return self.concept.name elif self.concept.path: @@ -241,20 +70,11 @@ def __get_display_name(self): else: return "" - def configure_element(self): - self.name_label.configure(text=self.__get_display_name()) - self.image.configure(light_image=self.__get_preview_image()) - try: - if hasattr(self.concept, '_search_cache'): - delattr(self.concept, '_search_cache') - except AttributeError: - pass - - def __get_preview_image(self): + def _get_preview_image(self): preview_path = "resources/icons/icon.png" glob_pattern = "**/*.*" if getattr(self.concept, 'include_subdirectories', False) else "*.*" - concept_path = ConceptWindow.get_concept_path(getattr(self.concept, 'path', None)) + concept_path = ConceptWindowController.get_concept_path(getattr(self.concept, 'path', None)) if concept_path: for path in pathlib.Path(concept_path).glob(glob_pattern): if any(part.startswith('.') for part in path.relative_to(concept_path).parent.parts): @@ -268,7 +88,7 @@ def __get_preview_image(self): break try: image = load_image(preview_path, convert_mode="RGBA") - except (OSError): + except OSError: image = Image.new("RGBA", (150, 150), (200, 200, 200, 255)) size = min(image.width, image.height) image = image.crop(( @@ -279,8 +99,9 @@ def __get_preview_image(self): )) return image.resize((150, 150), Image.Resampling.BILINEAR) - def place_in_list(self): - index = getattr(self, 'visible_index', self.i) - x = index % 6 - y = index // 6 - self.grid(row=y, column=x, pady=5, padx=5) + def _clear_search_cache(self): + try: + if hasattr(self.concept, '_search_cache'): + delattr(self.concept, '_search_cache') + except AttributeError: + pass diff --git a/modules/ui/BaseConceptWindowView.py b/modules/ui/BaseConceptWindowView.py index f58879d5f..0f851f5fb 100644 --- a/modules/ui/BaseConceptWindowView.py +++ b/modules/ui/BaseConceptWindowView.py @@ -1,840 +1,405 @@ import fractions import math -import os -import pathlib -import platform -import random -import threading -import time -import traceback - -from modules.util import concept_stats, path_util -from modules.util.config.ConceptConfig import ConceptConfig -from modules.util.config.TrainConfig import TrainConfig + +from modules.util import path_util from modules.util.enum.BalancingStrategy import BalancingStrategy from modules.util.enum.ConceptType import ConceptType -from modules.util.image_util import load_image -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState - -from mgds.LoadingPipeline import LoadingPipeline -from mgds.OutputPipelineModule import OutputPipelineModule -from mgds.PipelineModule import PipelineModule -from mgds.pipelineModules.CapitalizeTags import CapitalizeTags -from mgds.pipelineModules.DropTags import DropTags -from mgds.pipelineModules.RandomBrightness import RandomBrightness -from mgds.pipelineModules.RandomCircularMaskShrink import ( - RandomCircularMaskShrink, -) -from mgds.pipelineModules.RandomContrast import RandomContrast -from mgds.pipelineModules.RandomFlip import RandomFlip -from mgds.pipelineModules.RandomHue import RandomHue -from mgds.pipelineModules.RandomMaskRotateCrop import RandomMaskRotateCrop -from mgds.pipelineModules.RandomRotate import RandomRotate -from mgds.pipelineModules.RandomSaturation import RandomSaturation -from mgds.pipelineModules.ShuffleTags import ShuffleTags -from mgds.pipelineModuleTypes.RandomAccessPipelineModule import ( - RandomAccessPipelineModule, -) - -import torch -from torchvision.transforms import functional - -import customtkinter as ctk -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 - - -class InputPipelineModule( - PipelineModule, - RandomAccessPipelineModule, -): - def __init__(self, data: dict): - super().__init__() - self.data = data - - def length(self) -> int: - return 1 - - def get_inputs(self) -> list[str]: - return [] - - def get_outputs(self) -> list[str]: - return list(self.data.keys()) - - def get_item(self, variation: int, index: int, requested_name: str = None) -> dict: - return self.data - - -class ConceptWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - train_config: TrainConfig, - concept: ConceptConfig, - ui_state: UIState, - image_ui_state: UIState, - text_ui_state: UIState, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) - - self.train_config = train_config - - self.concept = concept - self.ui_state = ui_state - self.image_ui_state = image_ui_state - 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.title("Concept") - self.geometry("800x700") - self.resizable(True, True) - - self.grid_rowconfigure(0, weight=1) - self.grid_columnconfigure(0, weight=1) - - tabview = ctk.CTkTabview(self) - tabview.grid(row=0, column=0, sticky="nsew") - - self.general_tab = self.__general_tab(tabview.add("general"), concept) - self.image_augmentation_tab = self.__image_augmentation_tab(tabview.add("image augmentation")) - self.text_augmentation_tab = self.__text_augmentation_tab(tabview.add("text augmentation")) - self.concept_stats_tab = self.__concept_stats_tab(tabview.add("statistics")) - - #automatic concept scan - self.scan_thread = threading.Thread(target=self.__auto_update_concept_stats, daemon=True) - self.scan_thread.start() - - components.button(self, 1, 0, "ok", self.__ok) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def __general_tab(self, master, concept: ConceptConfig): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, weight=1) + +class BaseConceptWindowView: + def __init__(self, components): + self.components = components + + def build_general_tab(self, frame, controller, ui_state, text_ui_state): # name - components.label(frame, 0, 0, "Name", + self.components.label(frame, 0, 0, "Name", tooltip="Name of the concept") - components.entry(frame, 0, 1, self.ui_state, "name") + self.components.entry(frame, 0, 1, ui_state, "name") # enabled - components.label(frame, 1, 0, "Enabled", + self.components.label(frame, 1, 0, "Enabled", tooltip="Enable or disable this concept") - components.switch(frame, 1, 1, self.ui_state, "enabled") + self.components.switch(frame, 1, 1, ui_state, "enabled") # concept type - components.label(frame, 2, 0, "Concept Type", + self.components.label(frame, 2, 0, "Concept Type", tooltip="STANDARD: Standard finetuning with the sample as training target\n" "VALIDATION: Use concept for validation instead of training\n" "PRIOR_PREDICTION: Use the sample to make a prediction using the model as it was before training. This prediction is then used as the training target " "for the model in training. This can be used as regularisation and to preserve prior model knowledge while finetuning the model on other concepts. " "Only implemented for LoRA.", wide_tooltip=True) - components.options(frame, 2, 1, [str(x) for x in list(ConceptType)], self.ui_state, "type") + self.components.options(frame, 2, 1, [str(x) for x in list(ConceptType)], ui_state, "type") # path - components.label(frame, 3, 0, "Path", + self.components.label(frame, 3, 0, "Path", tooltip="Path where the training data is located") - components.path_entry(frame, 3, 1, self.ui_state, "path", mode="dir") - components.button(frame, 3, 2, text="download now", command=self.__download_dataset_threaded, + self.components.path_entry(frame, 3, 1, ui_state, "path", mode="dir") + self.components.button(frame, 3, 2, text="download now", command=controller.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", + self.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.path_entry(frame, 4, 2, self.text_ui_state, "prompt_path", mode="file") + prompt_path_entry = self.components.path_entry(frame, 4, 2, text_ui_state, "prompt_path", mode="file") def set_prompt_path_entry_enabled(option: str): - if option == 'concept': - for child in prompt_path_entry.children.values(): - child.configure(state="normal") - else: - for child in prompt_path_entry.children.values(): - child.configure(state="disabled") + self.components.set_widget_enabled(prompt_path_entry, option == 'concept') - components.options_kv(frame, 4, 1, [ + self.components.options_kv(frame, 4, 1, [ ("From text file per sample", 'sample'), ("From single text file", 'concept'), ("From image file name", 'filename'), - ], self.text_ui_state, "prompt_source", command=set_prompt_path_entry_enabled) - set_prompt_path_entry_enabled(concept.text.prompt_source) + ], text_ui_state, "prompt_source", command=set_prompt_path_entry_enabled) + set_prompt_path_entry_enabled(controller.concept.text.prompt_source) # include subdirectories - components.label(frame, 5, 0, "Include Subdirectories", + self.components.label(frame, 5, 0, "Include Subdirectories", tooltip="Includes images from subdirectories into the dataset") - components.switch(frame, 5, 1, self.ui_state, "include_subdirectories") + self.components.switch(frame, 5, 1, ui_state, "include_subdirectories") # image variations - components.label(frame, 6, 0, "Image Variations", + self.components.label(frame, 6, 0, "Image Variations", tooltip="The number of different image versions to cache if latent caching is enabled.") - components.entry(frame, 6, 1, self.ui_state, "image_variations") + self.components.entry(frame, 6, 1, ui_state, "image_variations") # text variations - components.label(frame, 7, 0, "Text Variations", + self.components.label(frame, 7, 0, "Text Variations", tooltip="The number of different text versions to cache if latent caching is enabled.") - components.entry(frame, 7, 1, self.ui_state, "text_variations") + self.components.entry(frame, 7, 1, ui_state, "text_variations") # balancing - components.label(frame, 8, 0, "Balancing", + self.components.label(frame, 8, 0, "Balancing", tooltip="The number of samples used during training. Use repeats to multiply the concept, or samples to specify an exact number of samples used in each epoch.") - components.entry(frame, 8, 1, self.ui_state, "balancing") - components.options(frame, 8, 2, [str(x) for x in list(BalancingStrategy)], self.ui_state, "balancing_strategy") + self.components.entry(frame, 8, 1, ui_state, "balancing") + self.components.options(frame, 8, 2, [str(x) for x in list(BalancingStrategy)], ui_state, "balancing_strategy") # loss weight - components.label(frame, 9, 0, "Loss Weight", + self.components.label(frame, 9, 0, "Loss Weight", tooltip="The loss multiplyer for this concept.") - components.entry(frame, 9, 1, self.ui_state, "loss_weight") - - frame.pack(fill="both", expand=1) - return frame - - def __image_augmentation_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) + self.components.entry(frame, 9, 1, ui_state, "loss_weight") + def build_image_augmentation_tab(self, frame, controller, image_ui_state): # header - components.label(frame, 0, 1, "Random", + self.components.label(frame, 0, 1, "Random", tooltip="Enable this augmentation with random values") - components.label(frame, 0, 2, "Fixed", + self.components.label(frame, 0, 2, "Fixed", tooltip="Enable this augmentation with fixed values") # crop jitter - components.label(frame, 1, 0, "Crop Jitter", + self.components.label(frame, 1, 0, "Crop Jitter", tooltip="Enables random cropping of samples") - components.switch(frame, 1, 1, self.image_ui_state, "enable_crop_jitter") + self.components.switch(frame, 1, 1, image_ui_state, "enable_crop_jitter") # random flip - components.label(frame, 2, 0, "Random Flip", + self.components.label(frame, 2, 0, "Random Flip", tooltip="Randomly flip the sample during training") - components.switch(frame, 2, 1, self.image_ui_state, "enable_random_flip") - components.switch(frame, 2, 2, self.image_ui_state, "enable_fixed_flip") + self.components.switch(frame, 2, 1, image_ui_state, "enable_random_flip") + self.components.switch(frame, 2, 2, image_ui_state, "enable_fixed_flip") # random rotation - components.label(frame, 3, 0, "Random Rotation", + self.components.label(frame, 3, 0, "Random Rotation", tooltip="Randomly rotates the sample during training") - components.switch(frame, 3, 1, self.image_ui_state, "enable_random_rotate") - components.switch(frame, 3, 2, self.image_ui_state, "enable_fixed_rotate") - components.entry(frame, 3, 3, self.image_ui_state, "random_rotate_max_angle") + self.components.switch(frame, 3, 1, image_ui_state, "enable_random_rotate") + self.components.switch(frame, 3, 2, image_ui_state, "enable_fixed_rotate") + self.components.entry(frame, 3, 3, image_ui_state, "random_rotate_max_angle") # random brightness - components.label(frame, 4, 0, "Random Brightness", + self.components.label(frame, 4, 0, "Random Brightness", tooltip="Randomly adjusts the brightness of the sample during training") - components.switch(frame, 4, 1, self.image_ui_state, "enable_random_brightness") - components.switch(frame, 4, 2, self.image_ui_state, "enable_fixed_brightness") - components.entry(frame, 4, 3, self.image_ui_state, "random_brightness_max_strength") + self.components.switch(frame, 4, 1, image_ui_state, "enable_random_brightness") + self.components.switch(frame, 4, 2, image_ui_state, "enable_fixed_brightness") + self.components.entry(frame, 4, 3, image_ui_state, "random_brightness_max_strength") # random contrast - components.label(frame, 5, 0, "Random Contrast", + self.components.label(frame, 5, 0, "Random Contrast", tooltip="Randomly adjusts the contrast of the sample during training") - components.switch(frame, 5, 1, self.image_ui_state, "enable_random_contrast") - components.switch(frame, 5, 2, self.image_ui_state, "enable_fixed_contrast") - components.entry(frame, 5, 3, self.image_ui_state, "random_contrast_max_strength") + self.components.switch(frame, 5, 1, image_ui_state, "enable_random_contrast") + self.components.switch(frame, 5, 2, image_ui_state, "enable_fixed_contrast") + self.components.entry(frame, 5, 3, image_ui_state, "random_contrast_max_strength") # random saturation - components.label(frame, 6, 0, "Random Saturation", + self.components.label(frame, 6, 0, "Random Saturation", tooltip="Randomly adjusts the saturation of the sample during training") - components.switch(frame, 6, 1, self.image_ui_state, "enable_random_saturation") - components.switch(frame, 6, 2, self.image_ui_state, "enable_fixed_saturation") - components.entry(frame, 6, 3, self.image_ui_state, "random_saturation_max_strength") + self.components.switch(frame, 6, 1, image_ui_state, "enable_random_saturation") + self.components.switch(frame, 6, 2, image_ui_state, "enable_fixed_saturation") + self.components.entry(frame, 6, 3, image_ui_state, "random_saturation_max_strength") # random hue - components.label(frame, 7, 0, "Random Hue", + self.components.label(frame, 7, 0, "Random Hue", tooltip="Randomly adjusts the hue of the sample during training") - components.switch(frame, 7, 1, self.image_ui_state, "enable_random_hue") - components.switch(frame, 7, 2, self.image_ui_state, "enable_fixed_hue") - components.entry(frame, 7, 3, self.image_ui_state, "random_hue_max_strength") + self.components.switch(frame, 7, 1, image_ui_state, "enable_random_hue") + self.components.switch(frame, 7, 2, image_ui_state, "enable_fixed_hue") + self.components.entry(frame, 7, 3, image_ui_state, "random_hue_max_strength") # random circular mask shrink - components.label(frame, 8, 0, "Circular Mask Generation", + self.components.label(frame, 8, 0, "Circular Mask Generation", tooltip="Automatically create circular masks for masked training") - components.switch(frame, 8, 1, self.image_ui_state, "enable_random_circular_mask_shrink") + self.components.switch(frame, 8, 1, image_ui_state, "enable_random_circular_mask_shrink") # random rotate and crop - components.label(frame, 9, 0, "Random Rotate and Crop", + self.components.label(frame, 9, 0, "Random Rotate and Crop", tooltip="Randomly rotate the training samples and crop to the masked region") - components.switch(frame, 9, 1, self.image_ui_state, "enable_random_mask_rotate_crop") + self.components.switch(frame, 9, 1, image_ui_state, "enable_random_mask_rotate_crop") # circular mask generation - components.label(frame, 10, 0, "Resolution Override", + self.components.label(frame, 10, 0, "Resolution Override", tooltip="Override the resolution for this concept. Optionally specify multiple resolutions separated by a comma, or a single exact resolution in the format x") - components.switch(frame, 10, 2, self.image_ui_state, "enable_resolution_override") - components.entry(frame, 10, 3, self.image_ui_state, "resolution_override") - - # image - image_preview, filename_preview, caption_preview = self.__get_preview_image() - self.image = ctk.CTkImage( - light_image=image_preview, - size=image_preview.size, - ) - image_label = ctk.CTkLabel(master=frame, text="", image=self.image, height=300, width=300) - image_label.grid(row=0, column=4, rowspan=6) - - # refresh preview - update_button_frame = ctk.CTkFrame(master=frame, corner_radius=0, fg_color="transparent") - update_button_frame.grid(row=6, column=4, rowspan=6, sticky="nsew") - update_button_frame.grid_columnconfigure(1, weight=1) - - prev_preview_button = components.button(update_button_frame, 0, 0, "<", command=self.__prev_image_preview) - components.button(update_button_frame, 0, 1, "Update Preview", command=self.__update_image_preview) - next_preview_button = components.button(update_button_frame, 0, 2, ">", command=self.__next_image_preview) - preview_augmentations_switch = ctk.CTkSwitch(update_button_frame, text="Show Augmentations", variable=self.preview_augmentations, command=self.__update_image_preview) - preview_augmentations_switch.grid(row=1, column=0, columnspan=3, padx=5, pady=5) - - prev_preview_button.configure(width=40) - next_preview_button.configure(width=40) - - #caption and filename preview - self.filename_preview = ctk.CTkLabel(master=update_button_frame, text=filename_preview, width=300, anchor="nw", justify="left", padx=10, wraplength=280) - self.filename_preview.grid(row=2, column=0, columnspan=3) - self.caption_preview = ctk.CTkTextbox(master=update_button_frame, width = 300, height = 150, wrap="word", border_width=2) - self.caption_preview.insert(index="1.0", text=caption_preview) - self.caption_preview.configure(state="disabled") - self.caption_preview.grid(row=3, column=0, columnspan=3, rowspan=3) - - frame.pack(fill="both", expand=1) - return frame - - def __text_augmentation_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) + self.components.switch(frame, 10, 2, image_ui_state, "enable_resolution_override") + self.components.entry(frame, 10, 3, image_ui_state, "resolution_override") + def build_text_augmentation_tab(self, frame, controller, text_ui_state): # tag shuffling - components.label(frame, 0, 0, "Tag Shuffling", + self.components.label(frame, 0, 0, "Tag Shuffling", tooltip="Enables tag shuffling") - components.switch(frame, 0, 1, self.text_ui_state, "enable_tag_shuffling") + self.components.switch(frame, 0, 1, text_ui_state, "enable_tag_shuffling") # keep tag count - components.label(frame, 1, 0, "Tag Delimiter", + self.components.label(frame, 1, 0, "Tag Delimiter", tooltip="The delimiter between tags") - components.entry(frame, 1, 1, self.text_ui_state, "tag_delimiter") + self.components.entry(frame, 1, 1, text_ui_state, "tag_delimiter") # keep tag count - components.label(frame, 2, 0, "Keep Tag Count", + self.components.label(frame, 2, 0, "Keep Tag Count", tooltip="The number of tags at the start of the caption that are not shuffled or dropped") - components.entry(frame, 2, 1, self.text_ui_state, "keep_tags_count") + self.components.entry(frame, 2, 1, text_ui_state, "keep_tags_count") # tag dropout - components.label(frame, 3, 0, "Tag Dropout", + self.components.label(frame, 3, 0, "Tag Dropout", tooltip="Enables random dropout for tags in the captions.") - components.switch(frame, 3, 1, self.text_ui_state, "tag_dropout_enable") - components.label(frame, 4, 0, "Dropout Mode", + self.components.switch(frame, 3, 1, text_ui_state, "tag_dropout_enable") + self.components.label(frame, 4, 0, "Dropout Mode", tooltip="Method used to drop captions. 'Full' will drop the entire caption past the 'kept' tags with a certain probability, 'Random' will drop individual tags with the set probability, and 'Random Weighted' will linearly increase the probability of dropping tags, more likely to preseve tags near the front with full probability to drop at the end.") - components.options_kv(frame, 4, 1, [ + self.components.options_kv(frame, 4, 1, [ ("Full", 'FULL'), ("Random", 'RANDOM'), ("Random Weighted", 'RANDOM WEIGHTED'), - ], self.text_ui_state, "tag_dropout_mode", None) - components.label(frame, 4, 2, "Probability", + ], text_ui_state, "tag_dropout_mode", None) + self.components.label(frame, 4, 2, "Probability", tooltip="Probability to drop tags, from 0 to 1.") - components.entry(frame, 4, 3, self.text_ui_state, "tag_dropout_probability") + self.components.entry(frame, 4, 3, text_ui_state, "tag_dropout_probability") - components.label(frame, 5, 0, "Special Dropout Tags", + self.components.label(frame, 5, 0, "Special Dropout Tags", tooltip="List of tags which will be whitelisted/blacklisted by dropout. 'Whitelist' tags will never be dropped but all others may be, 'Blacklist' tags may be dropped but all others will never be, 'None' may drop any tags. Can specify either a delimiter-separated list in the field, or a file path to a .txt or .csv file with entries separated by newlines.") - components.options_kv(frame, 5, 1, [ + self.components.options_kv(frame, 5, 1, [ ("None", 'NONE'), ("Blacklist", 'BLACKLIST'), ("Whitelist", 'WHITELIST'), - ], self.text_ui_state, "tag_dropout_special_tags_mode", None) - components.entry(frame, 5, 2, self.text_ui_state, "tag_dropout_special_tags") - components.label(frame, 6, 0, "Special Tags Regex", + ], text_ui_state, "tag_dropout_special_tags_mode", None) + self.components.entry(frame, 5, 2, text_ui_state, "tag_dropout_special_tags") + self.components.label(frame, 6, 0, "Special Tags Regex", tooltip="Interpret special tags with regex, such as 'photo.*' to match 'photo, photograph, photon' but not 'telephoto'. Includes exception for '/(' and '/)' syntax found in many booru/e6 tags.") - components.switch(frame, 6, 1, self.text_ui_state, "tag_dropout_special_tags_regex") + self.components.switch(frame, 6, 1, text_ui_state, "tag_dropout_special_tags_regex") #capitalization randomization - components.label(frame, 7, 0, "Randomize Capitalization", + self.components.label(frame, 7, 0, "Randomize Capitalization", tooltip="Enables randomization of capitalization for tags in the caption.") - components.switch(frame, 7, 1, self.text_ui_state, "caps_randomize_enable") - components.label(frame, 7, 2, "Force Lowercase", + self.components.switch(frame, 7, 1, text_ui_state, "caps_randomize_enable") + self.components.label(frame, 7, 2, "Force Lowercase", tooltip="If enabled, converts the caption to lowercase before any further processing.") - components.switch(frame, 7, 3, self.text_ui_state, "caps_randomize_lowercase") + self.components.switch(frame, 7, 3, text_ui_state, "caps_randomize_lowercase") - components.label(frame, 8, 0, "Captialization Mode", + self.components.label(frame, 8, 0, "Captialization Mode", tooltip="Comma-separated list of types of capitalization randomization to perform. 'capslock' for ALL CAPS, 'title' for First Letter Of Every Word, 'first' for First word only, 'random' for rAndOMiZeD lEtTERs.") - components.entry(frame, 8, 1, self.text_ui_state, "caps_randomize_mode") - components.label(frame, 8, 2, "Probability", + self.components.entry(frame, 8, 1, text_ui_state, "caps_randomize_mode") + self.components.label(frame, 8, 2, "Probability", tooltip="Probability to randomize capitialization of each tag, from 0 to 1.") - components.entry(frame, 8, 3, self.text_ui_state, "caps_randomize_probability") - - frame.pack(fill="both", expand=1) - return frame + self.components.entry(frame, 8, 3, text_ui_state, "caps_randomize_probability") - def __concept_stats_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=150) - frame.grid_columnconfigure(1, weight=0, minsize=150) - frame.grid_columnconfigure(2, weight=0, minsize=150) - frame.grid_columnconfigure(3, weight=0, minsize=150) - - self.cancel_scan_flag = threading.Event() + def build_concept_stats_tab(self, frame, controller): + self.concept_stats_tab = frame #file size - self.file_size_label = components.label(frame, 1, 0, "Total Size", pad=0, - tooltip="Total size of all image, mask, and caption files in MB") - self.file_size_label.configure(font=ctk.CTkFont(underline=True)) - self.file_size_preview = components.label(frame, 2, 0, pad=0, text="-") + self.file_size_label = self.components.label(frame, 1, 0, "Total Size", pad=0, + tooltip="Total size of all image, mask, and caption files in MB", underline=True) + self.file_size_preview = self.components.label(frame, 2, 0, pad=0, text="-") #subdirectory count - self.dir_count_label = components.label(frame, 1, 1, "Directories", pad=0, - tooltip="Total number of directories including and under (if 'include subdirectories' is enabled) the main concept directory") - self.dir_count_label.configure(font=ctk.CTkFont(underline=True)) - self.dir_count_preview = components.label(frame, 2, 1, pad=0, text="-") + self.dir_count_label = self.components.label(frame, 1, 1, "Directories", pad=0, + tooltip="Total number of directories including and under (if 'include subdirectories' is enabled) the main concept directory", underline=True) + self.dir_count_preview = self.components.label(frame, 2, 1, pad=0, text="-") #basic img/vid stats - count of each type in the concept #the \n at the start of the label gives it better vertical spacing with other rows - self.image_count_label = components.label(frame, 3, 0, "\nTotal Images", pad=0, - tooltip="Total number of image files, any of the extensions " + str(path_util.SUPPORTED_IMAGE_EXTENSIONS) + ", excluding '-masklabel.png and -condlabel.png'") - self.image_count_label.configure(font=ctk.CTkFont(underline=True)) - self.image_count_preview = components.label(frame, 4, 0, pad=0, text="-") - self.video_count_label = components.label(frame, 3, 1, "\nTotal Videos", pad=0, - tooltip="Total number of video files, any of the extensions " + str(path_util.SUPPORTED_VIDEO_EXTENSIONS)) - self.video_count_label.configure(font=ctk.CTkFont(underline=True)) - self.video_count_preview = components.label(frame, 4, 1, pad=0, text="-") - self.mask_count_label = components.label(frame, 3, 2, "\nTotal Masks", pad=0, - tooltip="Total number of mask files, any file ending in '-masklabel.png'") - self.mask_count_label.configure(font=ctk.CTkFont(underline=True)) - self.mask_count_preview = components.label(frame, 4, 2, pad=0, text="-") - self.caption_count_label = components.label(frame, 3, 3, "\nTotal Captions", pad=0, - tooltip="Total number of caption files, any .txt file. With advanced scan, includes the total number of captions on separate lines across all files in parentheses.") - self.caption_count_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_count_preview = components.label(frame, 4, 3, pad=0, text="-") + self.image_count_label = self.components.label(frame, 3, 0, "\nTotal Images", pad=0, + tooltip="Total number of image files, any of the extensions " + str(path_util.SUPPORTED_IMAGE_EXTENSIONS) + ", excluding '-masklabel.png and -condlabel.png'", underline=True) + self.image_count_preview = self.components.label(frame, 4, 0, pad=0, text="-") + self.video_count_label = self.components.label(frame, 3, 1, "\nTotal Videos", pad=0, + tooltip="Total number of video files, any of the extensions " + str(path_util.SUPPORTED_VIDEO_EXTENSIONS), underline=True) + self.video_count_preview = self.components.label(frame, 4, 1, pad=0, text="-") + self.mask_count_label = self.components.label(frame, 3, 2, "\nTotal Masks", pad=0, + tooltip="Total number of mask files, any file ending in '-masklabel.png'", underline=True) + self.mask_count_preview = self.components.label(frame, 4, 2, pad=0, text="-") + self.caption_count_label = self.components.label(frame, 3, 3, "\nTotal Captions", pad=0, + tooltip="Total number of caption files, any .txt file. With advanced scan, includes the total number of captions on separate lines across all files in parentheses.", underline=True) + self.caption_count_preview = self.components.label(frame, 4, 3, pad=0, text="-") #advanced img/vid stats - how many img/vid files have a mask or caption of the same name - self.image_count_mask_label = components.label(frame, 5, 0, "\nImages with Masks", pad=0, - tooltip="Total number of image files with an associated mask") - self.image_count_mask_label.configure(font=ctk.CTkFont(underline=True)) - self.image_count_mask_preview = components.label(frame, 6, 0, pad=0, text="-") - self.mask_count_label_unpaired = components.label(frame, 5, 1, "\nUnpaired Masks", pad=0, - tooltip="Total number of mask files which lack a corresponding image file - if >0, check your data set!") - self.mask_count_label_unpaired.configure(font=ctk.CTkFont(underline=True)) - self.mask_count_preview_unpaired = components.label(frame, 6, 1, pad=0, text="-") + self.image_count_mask_label = self.components.label(frame, 5, 0, "\nImages with Masks", pad=0, + tooltip="Total number of image files with an associated mask", underline=True) + self.image_count_mask_preview = self.components.label(frame, 6, 0, pad=0, text="-") + self.mask_count_label_unpaired = self.components.label(frame, 5, 1, "\nUnpaired Masks", pad=0, + tooltip="Total number of mask files which lack a corresponding image file - if >0, check your data set!", underline=True) + self.mask_count_preview_unpaired = self.components.label(frame, 6, 1, pad=0, text="-") #currently no masks for videos? - self.image_count_caption_label = components.label(frame, 7, 0, "\nImages with Captions", pad=0, - tooltip="Total number of image files with an associated caption") - self.image_count_caption_label.configure(font=ctk.CTkFont(underline=True)) - self.image_count_caption_preview = components.label(frame, 8, 0, pad=0, text="-") - self.video_count_caption_label = components.label(frame, 7, 1, "\nVideos with Captions", pad=0, - tooltip="Total number of video files with an associated caption") - self.video_count_caption_label.configure(font=ctk.CTkFont(underline=True)) - self.video_count_caption_preview = components.label(frame, 8, 1, pad=0, text="-") - self.caption_count_label_unpaired = components.label(frame, 7, 2, "\nUnpaired Captions", pad=0, - tooltip="Total number of caption files which lack a corresponding image file - if >0, check your data set! If using 'from file name' or 'from single text file' then this can be ignored.") - self.caption_count_label_unpaired.configure(font=ctk.CTkFont(underline=True)) - self.caption_count_preview_unpaired = components.label(frame, 8, 2, pad=0, text="-") + self.image_count_caption_label = self.components.label(frame, 7, 0, "\nImages with Captions", pad=0, + tooltip="Total number of image files with an associated caption", underline=True) + self.image_count_caption_preview = self.components.label(frame, 8, 0, pad=0, text="-") + self.video_count_caption_label = self.components.label(frame, 7, 1, "\nVideos with Captions", pad=0, + tooltip="Total number of video files with an associated caption", underline=True) + self.video_count_caption_preview = self.components.label(frame, 8, 1, pad=0, text="-") + self.caption_count_label_unpaired = self.components.label(frame, 7, 2, "\nUnpaired Captions", pad=0, + tooltip="Total number of caption files which lack a corresponding image file - if >0, check your data set! If using 'from file name' or 'from single text file' then this can be ignored.", underline=True) + self.caption_count_preview_unpaired = self.components.label(frame, 8, 2, pad=0, text="-") #resolution info - self.pixel_max_label = components.label(frame, 9, 0, "\nMax Pixels", pad=0, - tooltip="Largest image in the concept by number of pixels (width * height)") - self.pixel_max_label.configure(font=ctk.CTkFont(underline=True)) - self.pixel_max_preview = components.label(frame, 10, 0, pad=0, text="-", wraplength=150) - self.pixel_avg_label = components.label(frame, 9, 1, "\nAvg Pixels", pad=0, - tooltip="Average size of images in the concept by number of pixels (width * height)") - self.pixel_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.pixel_avg_preview = components.label(frame, 10, 1, pad=0, text="-", wraplength=150) - self.pixel_min_label = components.label(frame, 9, 2, "\nMin Pixels", pad=0, - tooltip="Smallest image in the concept by number of pixels (width * height)") - self.pixel_min_label.configure(font=ctk.CTkFont(underline=True)) - self.pixel_min_preview = components.label(frame, 10, 2, pad=0, text="-", wraplength=150) + self.pixel_max_label = self.components.label(frame, 9, 0, "\nMax Pixels", pad=0, + tooltip="Largest image in the concept by number of pixels (width * height)", underline=True) + self.pixel_max_preview = self.components.label(frame, 10, 0, pad=0, text="-", wraplength=150) + self.pixel_avg_label = self.components.label(frame, 9, 1, "\nAvg Pixels", pad=0, + tooltip="Average size of images in the concept by number of pixels (width * height)", underline=True) + self.pixel_avg_preview = self.components.label(frame, 10, 1, pad=0, text="-", wraplength=150) + self.pixel_min_label = self.components.label(frame, 9, 2, "\nMin Pixels", pad=0, + tooltip="Smallest image in the concept by number of pixels (width * height)", underline=True) + self.pixel_min_preview = self.components.label(frame, 10, 2, pad=0, text="-", wraplength=150) #video length info - self.length_max_label = components.label(frame, 11, 0, "\nMax Length", pad=0, - tooltip="Longest video in the concept by number of frames") - self.length_max_label.configure(font=ctk.CTkFont(underline=True)) - self.length_max_preview = components.label(frame, 12, 0, pad=0, text="-", wraplength=150) - self.length_avg_label = components.label(frame, 11, 1, "\nAvg Length", pad=0, - tooltip="Average length of videos in the concept by number of frames") - self.length_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.length_avg_preview = components.label(frame, 12, 1, pad=0, text="-", wraplength=150) - self.length_min_label = components.label(frame, 11, 2, "\nMin Length", pad=0, - tooltip="Shortest video in the concept by number of frames") - self.length_min_label.configure(font=ctk.CTkFont(underline=True)) - self.length_min_preview = components.label(frame, 12, 2, pad=0, text="-", wraplength=150) + self.length_max_label = self.components.label(frame, 11, 0, "\nMax Length", pad=0, + tooltip="Longest video in the concept by number of frames", underline=True) + self.length_max_preview = self.components.label(frame, 12, 0, pad=0, text="-", wraplength=150) + self.length_avg_label = self.components.label(frame, 11, 1, "\nAvg Length", pad=0, + tooltip="Average length of videos in the concept by number of frames", underline=True) + self.length_avg_preview = self.components.label(frame, 12, 1, pad=0, text="-", wraplength=150) + self.length_min_label = self.components.label(frame, 11, 2, "\nMin Length", pad=0, + tooltip="Shortest video in the concept by number of frames", underline=True) + self.length_min_preview = self.components.label(frame, 12, 2, pad=0, text="-", wraplength=150) #video fps info - self.fps_max_label = components.label(frame, 13, 0, "\nMax FPS", pad=0, - tooltip="Video in concept with highest fps") - self.fps_max_label.configure(font=ctk.CTkFont(underline=True)) - self.fps_max_preview = components.label(frame, 14, 0, pad=0, text="-", wraplength=150) - self.fps_avg_label = components.label(frame, 13, 1, "\nAvg FPS", pad=0, - tooltip="Average fps of videos in the concept") - self.fps_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.fps_avg_preview = components.label(frame, 14, 1, pad=0, text="-", wraplength=150) - self.fps_min_label = components.label(frame, 13, 2, "\nMin FPS", pad=0, - tooltip="Video in concept with the lowest fps") - self.fps_min_label.configure(font=ctk.CTkFont(underline=True)) - self.fps_min_preview = components.label(frame, 14, 2, pad=0, text="-", wraplength=150) + self.fps_max_label = self.components.label(frame, 13, 0, "\nMax FPS", pad=0, + tooltip="Video in concept with highest fps", underline=True) + self.fps_max_preview = self.components.label(frame, 14, 0, pad=0, text="-", wraplength=150) + self.fps_avg_label = self.components.label(frame, 13, 1, "\nAvg FPS", pad=0, + tooltip="Average fps of videos in the concept", underline=True) + self.fps_avg_preview = self.components.label(frame, 14, 1, pad=0, text="-", wraplength=150) + self.fps_min_label = self.components.label(frame, 13, 2, "\nMin FPS", pad=0, + tooltip="Video in concept with the lowest fps", underline=True) + self.fps_min_preview = self.components.label(frame, 14, 2, pad=0, text="-", wraplength=150) #caption info - self.caption_max_label = components.label(frame, 15, 0, "\nMax Caption Length", pad=0, - tooltip="Largest caption in concept by character count. For token count, assume ~2 tokens/word") - self.caption_max_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_max_preview = components.label(frame, 16, 0, pad=0, text="-", wraplength=150) - self.caption_avg_label = components.label(frame, 15, 1, "\nAvg Caption Length", pad=0, - tooltip="Average length of caption in concept by character count. For token count, assume ~2 tokens/word") - self.caption_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_avg_preview = components.label(frame, 16, 1, pad=0, text="-", wraplength=150) - self.caption_min_label = components.label(frame, 15, 2, "\nMin Caption Length", pad=0, - tooltip="Smallest caption in concept by character count. For token count, assume ~2 tokens/word") - self.caption_min_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_min_preview = components.label(frame, 16, 2, pad=0, text="-", wraplength=150) + self.caption_max_label = self.components.label(frame, 15, 0, "\nMax Caption Length", pad=0, + tooltip="Largest caption in concept by character count. For token count, assume ~2 tokens/word", underline=True) + self.caption_max_preview = self.components.label(frame, 16, 0, pad=0, text="-", wraplength=150) + self.caption_avg_label = self.components.label(frame, 15, 1, "\nAvg Caption Length", pad=0, + tooltip="Average length of caption in concept by character count. For token count, assume ~2 tokens/word", underline=True) + self.caption_avg_preview = self.components.label(frame, 16, 1, pad=0, text="-", wraplength=150) + self.caption_min_label = self.components.label(frame, 15, 2, "\nMin Caption Length", pad=0, + tooltip="Smallest caption in concept by character count. For token count, assume ~2 tokens/word", underline=True) + self.caption_min_preview = self.components.label(frame, 16, 2, pad=0, text="-", wraplength=150) #aspect bucket info - self.aspect_bucket_label = components.label(frame, 17, 0, "\nAspect Bucketing", pad=0, + self.aspect_bucket_label = self.components.label(frame, 17, 0, "\nAspect Bucketing", pad=0, tooltip="Graph of all possible buckets and the number of images in each one, defined as height/width. Buckets range from 0.25 (4:1 extremely wide) to 4 (1:4 extremely tall). \ - Images which don't match a bucket exactly are cropped to the nearest one.") - self.aspect_bucket_label.configure(font=ctk.CTkFont(underline=True)) - self.small_bucket_label = components.label(frame, 17, 1, "\nSmallest Buckets", pad=0, - tooltip="Image buckets with the least nonzero total images - if 'batch size' is larger than this, these images will be ignored during training! See the wiki for more details.") - self.small_bucket_label.configure(font=ctk.CTkFont(underline=True)) - self.small_bucket_preview = components.label(frame, 18, 1, pad=0, text="-") - - #aspect bucketing plot, mostly copied from timestep preview graph - appearance_mode = AppearanceModeTracker.get_mode() - background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) - text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) - background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" - self.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 - - assert self.bucket_fig is None - self.bucket_fig, self.bucket_ax = plt.subplots(figsize=(7,3)) - self.canvas = FigureCanvasTkAgg(self.bucket_fig, master=frame) - self.canvas.get_tk_widget().grid(row=19, column=0, columnspan=4, rowspan=2) - self.bucket_fig.tight_layout() - self.bucket_fig.subplots_adjust(bottom=0.15) - - self.bucket_fig.set_facecolor(background_color) - self.bucket_ax.set_facecolor(background_color) - self.bucket_ax.spines['bottom'].set_color(self.text_color) - self.bucket_ax.spines['left'].set_color(self.text_color) - self.bucket_ax.spines['top'].set_visible(False) - self.bucket_ax.spines['right'].set_color(self.text_color) - self.bucket_ax.tick_params(axis='x', colors=self.text_color, which="both") - self.bucket_ax.tick_params(axis='y', colors=self.text_color, which="both") - self.bucket_ax.xaxis.label.set_color(self.text_color) - self.bucket_ax.yaxis.label.set_color(self.text_color) + Images which don't match a bucket exactly are cropped to the nearest one.", underline=True) + self.small_bucket_label = self.components.label(frame, 17, 1, "\nSmallest Buckets", pad=0, + tooltip="Image buckets with the least nonzero total images - if 'batch size' is larger than this, these images will be ignored during training! See the wiki for more details.", underline=True) + self.small_bucket_preview = self.components.label(frame, 18, 1, pad=0, text="-") #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), + self.refresh_basic_stats_button = self.components.button(master=frame, row=0, column=0, text="Refresh Basic", command=lambda: controller.get_concept_stats_threaded(self, False, 9999), tooltip="Reload basic statistics for the concept directory") - self.refresh_advanced_stats_button = components.button(master=frame, row=0, column=1, text="Refresh Advanced", command=lambda: self.__get_concept_stats_threaded(True, 9999), + self.refresh_advanced_stats_button = self.components.button(master=frame, row=0, column=1, text="Refresh Advanced", command=lambda: controller.get_concept_stats_threaded(self, True, 9999), tooltip="Reload advanced statistics for the concept directory") #run "basic" scan first before "advanced", seems to help the system cache the directories and run faster - self.cancel_stats_button = components.button(master=frame, row=0, column=2, text="Abort Scan", command=lambda: self.__cancel_concept_stats(), + self.cancel_stats_button = self.components.button(master=frame, row=0, column=2, text="Abort Scan", command=lambda: self._cancel_concept_stats(controller), tooltip="Stop the currently running scan if it's taking a long time - advanced scan will be slow on large folders and on HDDs") - self.processing_time = components.label(frame, 0, 3, text="-", tooltip="Time taken to process concept directory") - - frame.pack(fill="both", expand=1) - return frame - - def __prev_image_preview(self): - self.image_preview_file_index = max(self.image_preview_file_index - 1, 0) - self.__update_image_preview() - - def __next_image_preview(self): - self.image_preview_file_index += 1 - self.__update_image_preview() - - def __update_image_preview(self): - image_preview, filename_preview, caption_preview = self.__get_preview_image() - self.image.configure(light_image=image_preview, size=image_preview.size) - self.filename_preview.configure(text=filename_preview) - self.caption_preview.configure(state="normal") - self.caption_preview.delete(index1="1.0", index2="end") - self.caption_preview.insert(index="1.0", text=caption_preview) - self.caption_preview.configure(state="disabled") - - @staticmethod - def get_concept_path(path: str) -> str | None: - if os.path.isdir(path): - return path - try: - #don't download, only check if available locally: - return huggingface_hub.snapshot_download(repo_id=path, repo_type="dataset", local_files_only=True) - except Exception: - return None - - def __download_dataset(self): - try: - huggingface_hub.login(token=self.train_config.secrets.huggingface_token, new_session=False) - huggingface_hub.snapshot_download(repo_id=self.concept.path, repo_type="dataset") - except Exception: - traceback.print_exc() - - def __download_dataset_threaded(self): - download_thread = threading.Thread(target=self.__download_dataset, daemon=True) - download_thread.start() - - def _read_text_file_for_preview(self, file_path: str) -> str: - empty_msg = "[Empty prompt]" - try: - with open(file_path, "r") as f: - if self.preview_augmentations.get(): - lines = [line.strip() for line in f if line.strip()] - return random.choice(lines) if lines else empty_msg - content = f.read().strip() - return content if content else empty_msg - except FileNotFoundError: - return "File not found, please check the path" - except IsADirectoryError: - return "[Provided path is a directory, please correct the caption path]" - except PermissionError: - if platform.system() == "Windows": - return "[Permission denied, please check the file permissions or Windows Defender settings]" - else: - return "[Permission denied, please check the file permissions]" - except UnicodeDecodeError: - return "[Invalid file encoding. This should not happen, please report this issue]" - - def __get_preview_image(self): - preview_image_path = "resources/icons/icon.png" - file_index = -1 - glob_pattern = "**/*.*" if self.concept.include_subdirectories else "*.*" - - concept_path = self.get_concept_path(self.concept.path) - if concept_path: - for path in pathlib.Path(concept_path).glob(glob_pattern): - if any(part.startswith('.') for part in path.relative_to(concept_path).parent.parts): - continue - extension = os.path.splitext(path)[1] - if path.is_file() and path_util.is_supported_image_extension(extension) \ - and not path.name.endswith("-masklabel.png") and not path.name.endswith("-condlabel.png"): - preview_image_path = path_util.canonical_join(concept_path, path) - file_index += 1 - if file_index == self.image_preview_file_index: - break - - image = load_image(preview_image_path, 'RGB') - image_tensor = functional.to_tensor(image) - - splitext = os.path.splitext(preview_image_path) - preview_mask_path = path_util.canonical_join(splitext[0] + "-masklabel.png") - if not os.path.isfile(preview_mask_path): - preview_mask_path = None - - if preview_mask_path: - mask = Image.open(preview_mask_path).convert("L") - mask_tensor = functional.to_tensor(mask) - else: - mask_tensor = torch.ones((1, image_tensor.shape[1], image_tensor.shape[2])) + self.processing_time = self.components.label(frame, 0, 3, text="-", tooltip="Time taken to process concept directory") - source = self.concept.text.prompt_source - preview_p = pathlib.Path(preview_image_path) - if source == "filename": - prompt_output = preview_p.stem or "[Empty prompt]" - else: - file_map = { - "sample": preview_p.with_suffix(".txt"), - "concept": pathlib.Path(self.concept.text.prompt_path) if self.concept.text.prompt_path else None, - } - file_path = file_map.get(source) - prompt_output = self._read_text_file_for_preview(str(file_path)) if file_path else "[Empty prompt]" - - modules = [] - if self.preview_augmentations.get(): - input_module = InputPipelineModule({ - 'true': True, - 'image': image_tensor, - 'mask': mask_tensor, - 'enable_random_flip': self.concept.image.enable_random_flip, - 'enable_fixed_flip': self.concept.image.enable_fixed_flip, - 'enable_random_rotate': self.concept.image.enable_random_rotate, - 'enable_fixed_rotate': self.concept.image.enable_fixed_rotate, - 'random_rotate_max_angle': self.concept.image.random_rotate_max_angle, - 'enable_random_brightness': self.concept.image.enable_random_brightness, - 'enable_fixed_brightness': self.concept.image.enable_fixed_brightness, - 'random_brightness_max_strength': self.concept.image.random_brightness_max_strength, - 'enable_random_contrast': self.concept.image.enable_random_contrast, - 'enable_fixed_contrast': self.concept.image.enable_fixed_contrast, - 'random_contrast_max_strength': self.concept.image.random_contrast_max_strength, - 'enable_random_saturation': self.concept.image.enable_random_saturation, - 'enable_fixed_saturation': self.concept.image.enable_fixed_saturation, - 'random_saturation_max_strength': self.concept.image.random_saturation_max_strength, - 'enable_random_hue': self.concept.image.enable_random_hue, - 'enable_fixed_hue': self.concept.image.enable_fixed_hue, - 'random_hue_max_strength': self.concept.image.random_hue_max_strength, - 'enable_random_circular_mask_shrink': self.concept.image.enable_random_circular_mask_shrink, - 'enable_random_mask_rotate_crop': self.concept.image.enable_random_mask_rotate_crop, - - 'prompt' : prompt_output, - 'tag_dropout_enable' : self.concept.text.tag_dropout_enable, - 'tag_dropout_probability' : self.concept.text.tag_dropout_probability, - 'tag_dropout_mode' : self.concept.text.tag_dropout_mode, - 'tag_dropout_special_tags' : self.concept.text.tag_dropout_special_tags, - 'tag_dropout_special_tags_mode' : self.concept.text.tag_dropout_special_tags_mode, - 'tag_delimiter' : self.concept.text.tag_delimiter, - 'keep_tags_count' : self.concept.text.keep_tags_count, - 'tag_dropout_special_tags_regex' : self.concept.text.tag_dropout_special_tags_regex, - 'caps_randomize_enable' : self.concept.text.caps_randomize_enable, - 'caps_randomize_probability' : self.concept.text.caps_randomize_probability, - 'caps_randomize_mode' : self.concept.text.caps_randomize_mode, - 'caps_randomize_lowercase' : self.concept.text.caps_randomize_lowercase, - 'enable_tag_shuffling' : self.concept.text.enable_tag_shuffling, - }) - - circular_mask_shrink = RandomCircularMaskShrink(mask_name='mask', shrink_probability=1.0, shrink_factor_min=0.2, shrink_factor_max=1.0, enabled_in_name='enable_random_circular_mask_shrink') - random_mask_rotate_crop = RandomMaskRotateCrop(mask_name='mask', additional_names=['image'], min_size=512, min_padding_percent=10, max_padding_percent=30, max_rotate_angle=20, enabled_in_name='enable_random_mask_rotate_crop') - random_flip = RandomFlip(names=['image', 'mask'], enabled_in_name='enable_random_flip', fixed_enabled_in_name='enable_fixed_flip') - random_rotate = RandomRotate(names=['image', 'mask'], enabled_in_name='enable_random_rotate', fixed_enabled_in_name='enable_fixed_rotate', max_angle_in_name='random_rotate_max_angle') - random_brightness = RandomBrightness(names=['image'], enabled_in_name='enable_random_brightness', fixed_enabled_in_name='enable_fixed_brightness', max_strength_in_name='random_brightness_max_strength') - random_contrast = RandomContrast(names=['image'], enabled_in_name='enable_random_contrast', fixed_enabled_in_name='enable_fixed_contrast', max_strength_in_name='random_contrast_max_strength') - random_saturation = RandomSaturation(names=['image'], enabled_in_name='enable_random_saturation', fixed_enabled_in_name='enable_fixed_saturation', max_strength_in_name='random_saturation_max_strength') - random_hue = RandomHue(names=['image'], enabled_in_name='enable_random_hue', fixed_enabled_in_name='enable_fixed_hue', max_strength_in_name='random_hue_max_strength') - drop_tags = DropTags(text_in_name='prompt', enabled_in_name='tag_dropout_enable', probability_in_name='tag_dropout_probability', dropout_mode_in_name='tag_dropout_mode', - special_tags_in_name='tag_dropout_special_tags', special_tag_mode_in_name='tag_dropout_special_tags_mode', delimiter_in_name='tag_delimiter', - keep_tags_count_in_name='keep_tags_count', text_out_name='prompt', regex_enabled_in_name='tag_dropout_special_tags_regex') - caps_randomize = CapitalizeTags(text_in_name='prompt', enabled_in_name='caps_randomize_enable', probability_in_name='caps_randomize_probability', - capitalize_mode_in_name='caps_randomize_mode', delimiter_in_name='tag_delimiter', convert_lowercase_in_name='caps_randomize_lowercase', text_out_name='prompt') - shuffle_tags = ShuffleTags(text_in_name='prompt', enabled_in_name='enable_tag_shuffling', delimiter_in_name='tag_delimiter', keep_tags_count_in_name='keep_tags_count', text_out_name='prompt') - output_module = OutputPipelineModule(['image', 'mask', 'prompt']) - - modules = [ - input_module, - circular_mask_shrink, - random_mask_rotate_crop, - random_flip, - random_rotate, - random_brightness, - random_contrast, - random_saturation, - random_hue, - drop_tags, - caps_randomize, - shuffle_tags, - output_module, - ] - - pipeline = LoadingPipeline( - device=torch.device('cpu'), - modules=modules, - batch_size=1, - seed=random.randint(0, 2**30), - state=None, - initial_epoch=0, - initial_index=0, - ) - - data = pipeline.__next__() - image_tensor = data['image'] - mask_tensor = data['mask'] - prompt_output = data['prompt'] - - filename_output = os.path.basename(preview_image_path) - - mask_tensor = torch.clamp(mask_tensor, 0.3, 1) - image_tensor = image_tensor * mask_tensor - - image = functional.to_pil_image(image_tensor) - - image.thumbnail((300, 300)) - - return image, filename_output, prompt_output - - def __update_concept_stats(self): + def _update_concept_stats(self, controller): #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") + self.components.set_label_text(self.file_size_preview, str(int(controller.concept.concept_stats["file_size"]/1048576)) + " MB") + self.components.set_label_text(self.processing_time, str(round(controller.concept.concept_stats["processing_time"], 2)) + " s") #directory count - self.dir_count_preview.configure(text=self.concept.concept_stats["directory_count"]) + self.components.set_label_text(self.dir_count_preview, controller.concept.concept_stats["directory_count"]) #image count - self.image_count_preview.configure(text=self.concept.concept_stats["image_count"]) - self.image_count_mask_preview.configure(text=self.concept.concept_stats["image_with_mask_count"]) - self.image_count_caption_preview.configure(text=self.concept.concept_stats["image_with_caption_count"]) + self.components.set_label_text(self.image_count_preview, controller.concept.concept_stats["image_count"]) + self.components.set_label_text(self.image_count_mask_preview, controller.concept.concept_stats["image_with_mask_count"]) + self.components.set_label_text(self.image_count_caption_preview, controller.concept.concept_stats["image_with_caption_count"]) #video count - self.video_count_preview.configure(text=self.concept.concept_stats["video_count"]) - #self.video_count_mask_preview.configure(text=self.concept.concept_stats["video_with_mask_count"]) - self.video_count_caption_preview.configure(text=self.concept.concept_stats["video_with_caption_count"]) + self.components.set_label_text(self.video_count_preview, controller.concept.concept_stats["video_count"]) + #self.components.set_label_text(self.video_count_mask_preview, controller.concept.concept_stats["video_with_mask_count"]) + self.components.set_label_text(self.video_count_caption_preview, controller.concept.concept_stats["video_with_caption_count"]) #mask count - self.mask_count_preview.configure(text=self.concept.concept_stats["mask_count"]) - self.mask_count_preview_unpaired.configure(text=self.concept.concept_stats["unpaired_masks"]) + self.components.set_label_text(self.mask_count_preview, controller.concept.concept_stats["mask_count"]) + self.components.set_label_text(self.mask_count_preview_unpaired, controller.concept.concept_stats["unpaired_masks"]) #caption count - if self.concept.concept_stats["subcaption_count"] > 0: - self.caption_count_preview.configure(text=f'{self.concept.concept_stats["caption_count"]} ({self.concept.concept_stats["subcaption_count"]})') + if controller.concept.concept_stats["subcaption_count"] > 0: + self.components.set_label_text(self.caption_count_preview, f'{controller.concept.concept_stats["caption_count"]} ({controller.concept.concept_stats["subcaption_count"]})') else: - self.caption_count_preview.configure(text=self.concept.concept_stats["caption_count"]) - self.caption_count_preview_unpaired.configure(text=self.concept.concept_stats["unpaired_captions"]) + self.components.set_label_text(self.caption_count_preview, controller.concept.concept_stats["caption_count"]) + self.components.set_label_text(self.caption_count_preview_unpaired, controller.concept.concept_stats["unpaired_captions"]) #resolution info - max_pixels = self.concept.concept_stats["max_pixels"] - avg_pixels = self.concept.concept_stats["avg_pixels"] - min_pixels = self.concept.concept_stats["min_pixels"] - - if any(isinstance(x, str) for x in [max_pixels, avg_pixels, min_pixels]) or self.concept.concept_stats["image_count"] == 0: #will be str if adv stats were not taken - self.pixel_max_preview.configure(text="-") - self.pixel_avg_preview.configure(text="-") - self.pixel_min_preview.configure(text="-") + max_pixels = controller.concept.concept_stats["max_pixels"] + avg_pixels = controller.concept.concept_stats["avg_pixels"] + min_pixels = controller.concept.concept_stats["min_pixels"] + + if any(isinstance(x, str) for x in [max_pixels, avg_pixels, min_pixels]) or controller.concept.concept_stats["image_count"] == 0: #will be str if adv stats were not taken + self.components.set_label_text(self.pixel_max_preview, "-") + self.components.set_label_text(self.pixel_avg_preview, "-") + self.components.set_label_text(self.pixel_min_preview, "-") else: #formatted as (#pixels/1000000) MP, width x height, \n filename - self.pixel_max_preview.configure(text=f'{str(round(max_pixels[0]/1000000, 2))} MP, {max_pixels[2]}\n{max_pixels[1]}') - self.pixel_avg_preview.configure(text=f'{str(round(avg_pixels/1000000, 2))} MP, ~{int(math.sqrt(avg_pixels))}w x {int(math.sqrt(avg_pixels))}h') - self.pixel_min_preview.configure(text=f'{str(round(min_pixels[0]/1000000, 2))} MP, {min_pixels[2]}\n{min_pixels[1]}') + self.components.set_label_text(self.pixel_max_preview, f'{str(round(max_pixels[0]/1000000, 2))} MP, {max_pixels[2]}\n{max_pixels[1]}') + self.components.set_label_text(self.pixel_avg_preview, f'{str(round(avg_pixels/1000000, 2))} MP, ~{int(math.sqrt(avg_pixels))}w x {int(math.sqrt(avg_pixels))}h') + self.components.set_label_text(self.pixel_min_preview, f'{str(round(min_pixels[0]/1000000, 2))} MP, {min_pixels[2]}\n{min_pixels[1]}') #video length and fps info - max_length = self.concept.concept_stats["max_length"] - avg_length = self.concept.concept_stats["avg_length"] - min_length = self.concept.concept_stats["min_length"] - max_fps = self.concept.concept_stats["max_fps"] - avg_fps = self.concept.concept_stats["avg_fps"] - min_fps = self.concept.concept_stats["min_fps"] - - if any(isinstance(x, str) for x in [max_length, avg_length, min_length]) or self.concept.concept_stats["video_count"] == 0: #will be str if adv stats were not taken - self.length_max_preview.configure(text="-") - self.length_avg_preview.configure(text="-") - self.length_min_preview.configure(text="-") - self.fps_max_preview.configure(text="-") - self.fps_avg_preview.configure(text="-") - self.fps_min_preview.configure(text="-") + max_length = controller.concept.concept_stats["max_length"] + avg_length = controller.concept.concept_stats["avg_length"] + min_length = controller.concept.concept_stats["min_length"] + max_fps = controller.concept.concept_stats["max_fps"] + avg_fps = controller.concept.concept_stats["avg_fps"] + min_fps = controller.concept.concept_stats["min_fps"] + + if any(isinstance(x, str) for x in [max_length, avg_length, min_length]) or controller.concept.concept_stats["video_count"] == 0: #will be str if adv stats were not taken + self.components.set_label_text(self.length_max_preview, "-") + self.components.set_label_text(self.length_avg_preview, "-") + self.components.set_label_text(self.length_min_preview, "-") + self.components.set_label_text(self.fps_max_preview, "-") + self.components.set_label_text(self.fps_avg_preview, "-") + self.components.set_label_text(self.fps_min_preview, "-") else: #formatted as (#frames) frames \n filename - self.length_max_preview.configure(text=f'{int(max_length[0])} frames\n{max_length[1]}') - self.length_avg_preview.configure(text=f'{int(avg_length)} frames') - self.length_min_preview.configure(text=f'{int(min_length[0])} frames\n{min_length[1]}') + self.components.set_label_text(self.length_max_preview, f'{int(max_length[0])} frames\n{max_length[1]}') + self.components.set_label_text(self.length_avg_preview, f'{int(avg_length)} frames') + self.components.set_label_text(self.length_min_preview, f'{int(min_length[0])} frames\n{min_length[1]}') #formatted as (#fps) fps \n filename - self.fps_max_preview.configure(text=f'{int(max_fps[0])} fps\n{max_fps[1]}') - self.fps_avg_preview.configure(text=f'{int(avg_fps)} fps') - self.fps_min_preview.configure(text=f'{int(min_fps[0])} fps\n{min_fps[1]}') + self.components.set_label_text(self.fps_max_preview, f'{int(max_fps[0])} fps\n{max_fps[1]}') + self.components.set_label_text(self.fps_avg_preview, f'{int(avg_fps)} fps') + self.components.set_label_text(self.fps_min_preview, f'{int(min_fps[0])} fps\n{min_fps[1]}') #caption info - max_caption_length = self.concept.concept_stats["max_caption_length"] - avg_caption_length = self.concept.concept_stats["avg_caption_length"] - min_caption_length = self.concept.concept_stats["min_caption_length"] - - if any(isinstance(x, str) for x in [max_caption_length, avg_caption_length, min_caption_length]) or self.concept.concept_stats["caption_count"] == 0: #will be str if adv stats were not taken - self.caption_max_preview.configure(text="-") - self.caption_avg_preview.configure(text="-") - self.caption_min_preview.configure(text="-") + max_caption_length = controller.concept.concept_stats["max_caption_length"] + avg_caption_length = controller.concept.concept_stats["avg_caption_length"] + min_caption_length = controller.concept.concept_stats["min_caption_length"] + + if any(isinstance(x, str) for x in [max_caption_length, avg_caption_length, min_caption_length]) or controller.concept.concept_stats["caption_count"] == 0: #will be str if adv stats were not taken + self.components.set_label_text(self.caption_max_preview, "-") + self.components.set_label_text(self.caption_avg_preview, "-") + self.components.set_label_text(self.caption_min_preview, "-") else: #formatted as (#chars) chars, (#words) words, \n filename - self.caption_max_preview.configure(text=f'{max_caption_length[0]} chars, {max_caption_length[2]} words\n{max_caption_length[1]}') - self.caption_avg_preview.configure(text=f'{int(avg_caption_length[0])} chars, {int(avg_caption_length[1])} words') - self.caption_min_preview.configure(text=f'{min_caption_length[0]} chars, {min_caption_length[2]} words\n{min_caption_length[1]}') + self.components.set_label_text(self.caption_max_preview, f'{max_caption_length[0]} chars, {max_caption_length[2]} words\n{max_caption_length[1]}') + self.components.set_label_text(self.caption_avg_preview, f'{int(avg_caption_length[0])} chars, {int(avg_caption_length[1])} words') + self.components.set_label_text(self.caption_min_preview, f'{min_caption_length[0]} chars, {min_caption_length[2]} words\n{min_caption_length[1]}') #aspect bucketing - aspect_buckets = self.concept.concept_stats["aspect_buckets"] + aspect_buckets = controller.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 @@ -846,7 +411,7 @@ def __update_concept_stats(self): 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) + self.components.set_label_text(self.small_bucket_preview, min_bucket_str) self.bucket_ax.cla() aspects = [str(x) for x in list(aspect_buckets.keys())] @@ -866,69 +431,13 @@ def decimal_to_aspect_ratio(self, value : float): aspect_string = f'{aspect_fraction.denominator}:{aspect_fraction.numerator}' return aspect_string - def __get_concept_stats(self, advanced_checks: bool, wait_time: float): - start_time = time.perf_counter() - last_update = time.perf_counter() - self.cancel_scan_flag.clear() - self.concept_stats_tab.after(0, self.__disable_scan_buttons) - concept_path = self.get_concept_path(self.concept.path) - - if not concept_path: - print(f"Unable to get statistics for concept path: {self.concept.path}") - self.concept_stats_tab.after(0, self.__enable_scan_buttons) - return - subfolders = [concept_path] - - 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: - 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 - subfolders.extend([f for f in os.scandir(path) if f.is_dir() and not f.name.startswith('.')]) - self.concept.concept_stats = stats_dict - #update GUI approx every half second - if time.perf_counter() > (last_update + 0.5): - 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) - - 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): - self.refresh_basic_stats_button.configure(state="disabled") - self.refresh_advanced_stats_button.configure(state="disabled") - - def __enable_scan_buttons(self): - self.refresh_basic_stats_button.configure(state="normal") - self.refresh_advanced_stats_button.configure(state="normal") - - def __cancel_concept_stats(self): - self.cancel_scan_flag.set() - - def __auto_update_concept_stats(self): - try: - self.__update_concept_stats() #load stats from config if available, else raises KeyError - if self.concept.concept_stats["file_size"] == 0: #force rescan if empty - raise KeyError - except KeyError: - concept_path = self.get_concept_path(self.concept.path) - if concept_path: - self.__get_concept_stats(False, 2) #force rescan if config is empty, timeout of 2 sec - if self.concept.concept_stats["processing_time"] < 0.1: - self.__get_concept_stats(True, 2) #do advanced scan automatically if basic took <0.1s - - def destroy(self): - if self.bucket_fig is not None: - plt.close(self.bucket_fig) - self.bucket_fig = None - - super().destroy() - - def __ok(self): - self.destroy() + def _disable_scan_buttons(self): + self.components.set_widget_enabled(self.refresh_basic_stats_button, False) + self.components.set_widget_enabled(self.refresh_advanced_stats_button, False) + + def _enable_scan_buttons(self): + self.components.set_widget_enabled(self.refresh_basic_stats_button, True) + self.components.set_widget_enabled(self.refresh_advanced_stats_button, True) + + def _cancel_concept_stats(self, controller): + controller.cancel_scan_flag.set() diff --git a/modules/ui/BaseConfigListView.py b/modules/ui/BaseConfigListView.py index 75d69252a..2165e854b 100644 --- a/modules/ui/BaseConfigListView.py +++ b/modules/ui/BaseConfigListView.py @@ -1,27 +1,63 @@ -import contextlib import copy import json import os -import tkinter as tk -from abc import ABCMeta, abstractmethod +from abc import ABC, abstractmethod from modules.util import path_util -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.UIState import UIState import customtkinter as ctk -class ConfigList(metaclass=ABCMeta): +class BaseConfigListView(ABC): - def __init__( + def __init__(self, components): + self.components = components + + @abstractmethod + def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): + pass + + @abstractmethod + def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: + pass + + @abstractmethod + def _create_top_frame(self, master): + pass + + @abstractmethod + def _create_element_list_frame(self, master): + pass + + @abstractmethod + def _remove_widget_from_layout(self, widget): + pass + + @abstractmethod + def _destroy_widget(self, widget): + pass + + @abstractmethod + def _destroy_frame(self, frame): + pass + + @abstractmethod + def _wait_for_window(self, window): + pass + + @abstractmethod + def _show_name_dialog(self, callback): + pass + + def _refresh_show_disabled_text(self): + return + + def build( self, master, - train_config: TrainConfig, - ui_state: UIState, + controller, + ui_state, from_external_file: bool, attr_name: str = "", enable_key: str = "enabled", @@ -29,11 +65,11 @@ def __init__( default_config_name: str = "", add_button_text: str = "", add_button_tooltip: str = "", - is_full_width: bool = "", + is_full_width: bool = False, show_toggle_button: bool = False, ): self.master = master - self.train_config = train_config + self.controller = controller self.ui_state = ui_state self.from_external_file = from_external_file self.attr_name = attr_name @@ -54,13 +90,8 @@ def __init__( self.is_opening_window = False self._is_current_item_enabled = False - self.master.grid_rowconfigure(0, weight=0) - self.master.grid_rowconfigure(1, weight=1) - self.master.grid_columnconfigure(0, weight=1) - if self.from_external_file: - self.top_frame = ctk.CTkFrame(self.master, fg_color="transparent") - self.top_frame.grid(row=0, column=0, sticky="nsew") + self.top_frame = self._create_top_frame(master) self.configs_dropdown = None self.element_list = None @@ -68,59 +99,27 @@ def __init__( self.configs = [] self.__load_available_config_names() - self.current_config = getattr(self.train_config, self.attr_name) + self.current_config = getattr(self.controller.train_config, self.attr_name) self.widgets = [] - self.__load_current_config(getattr(self.train_config, self.attr_name)) + self.__load_current_config(getattr(self.controller.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) + self.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) + self.components.button(self.top_frame, 0, 2, add_button_text, self.__add_element, tooltip=add_button_tooltip, width=30, padx=5) 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.top_frame = self._create_top_frame(master) + self.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.current_config = getattr(self.controller.train_config, self.attr_name) self.element_list = None 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.toggle_button = self.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() - - - @abstractmethod - def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - pass - - @abstractmethod - def create_new_element(self) -> BaseConfig: - pass - - @abstractmethod - def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: - pass - - def _refresh_show_disabled_text(self): - return - - def _reset_filters(self): # pragma: no cover - default noop - search_var = getattr(self, 'search_var', None) - filter_var = getattr(self, 'filter_var', None) - show_disabled_var = getattr(self, 'show_disabled_var', None) - - if search_var: - search_var.set("") - if filter_var: - filter_var.set("ALL") - if show_disabled_var: - show_disabled_var.set(True) - if search_var and hasattr(self, '_update_filters'): - self._update_filters() - def _update_item_enabled_state(self): # Only count items that match current filters self._is_current_item_enabled = any( @@ -154,14 +153,14 @@ def __create_configs_dropdown(self): if self.configs_dropdown is not None: self.configs_dropdown.destroy() - self.configs_dropdown = components.options_kv( + self.configs_dropdown = self.components.options_kv( self.top_frame, 0, 0, self.configs, self.ui_state, self.attr_name, self.__load_current_config ) self._update_toggle_button_text() def _create_element_list(self, **filters): if not self.from_external_file: - self.current_config = getattr(self.train_config, self.attr_name) + self.current_config = getattr(self.controller.train_config, self.attr_name) self.filters.update(filters) @@ -175,13 +174,9 @@ def _create_element_list(self, **filters): def _initialize_all_widgets(self): self.widgets = [] if self.element_list is not None: - self.element_list.destroy() + self._destroy_frame(self.element_list) - self.element_list = ctk.CTkScrollableFrame(self.master, fg_color="transparent") - self.element_list.grid(row=1, column=0, sticky="nsew") - - if self.is_full_width: - self.element_list.grid_columnconfigure(0, weight=1) + self.element_list = self._create_element_list_frame(self.master) for i, element in enumerate(self.current_config): widget = self.create_widget( @@ -205,7 +200,7 @@ def _update_widget_visibility(self): widget.place_in_list() visible_index += 1 else: - widget.grid_remove() + self._remove_widget_from_layout(widget) def __load_available_config_names(self): if os.path.isdir(self.config_dir): @@ -228,10 +223,10 @@ def __create_config(self, name: str): self.__create_configs_dropdown() def __add_config(self): - dialogs.StringInputDialog(self.master, "name", "Name", self.__create_config) + self._show_name_dialog(self.__create_config) def __add_element(self): - new_element = self.create_new_element() + new_element = self.controller.create_new_element() 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: @@ -276,8 +271,7 @@ 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() + self._destroy_widget(removed) # Reindex remaining widgets for idx, widget in enumerate(self.widgets): widget.i = idx @@ -294,7 +288,7 @@ def __load_current_config(self, filename): loaded_config_json = json.load(f) for element_json in loaded_config_json: - element = self.create_new_element().from_dict(element_json) + element = self.controller.create_new_element().from_dict(element_json) self.current_config.append(element) except (FileNotFoundError, json.JSONDecodeError) as e: print(f"Failed to load config from {filename}: {e}") @@ -315,7 +309,7 @@ def save_current_config(self): os.makedirs(self.config_dir, exist_ok=True) write_json_atomic( - getattr(self.train_config, self.attr_name), + getattr(self.controller.train_config, self.attr_name), [element.to_dict() for element in self.current_config] ) except (OSError) as e: @@ -324,10 +318,7 @@ def save_current_config(self): self._update_toggle_button_text() if self.widgets_initialized: - try: - self._update_widget_visibility() - except (tk.TclError, AttributeError) as e: - print.debug(f"Widget visibility update failed: {e}") + self._update_widget_visibility() # let subclass refresh any show-disabled UI if hasattr(self, '_refresh_show_disabled_text'): @@ -336,13 +327,27 @@ def save_current_config(self): def _element_matches_filters(self, element): return True # Show all by default + def _reset_filters(self): # pragma: no cover - default noop + search_var = getattr(self, 'search_var', None) + filter_var = getattr(self, 'filter_var', None) + show_disabled_var = getattr(self, 'show_disabled_var', None) + + if search_var: + search_var.set("") + if filter_var: + filter_var.set("ALL") + if show_disabled_var: + show_disabled_var.set(True) + if search_var and hasattr(self, '_update_filters'): + self._update_filters() + def __open_element_window(self, i, ui_state): if self.is_opening_window: return self.is_opening_window = True try: window = self.open_element_window(i, ui_state) - self.master.wait_window(window) + self._wait_for_window(window) try: if self.widgets is not None and 0 <= i < len(self.widgets): self.widgets[i].configure_element() diff --git a/modules/ui/BaseConvertModelUIView.py b/modules/ui/BaseConvertModelUIView.py index 6cb1b507a..7e2883b1e 100644 --- a/modules/ui/BaseConvertModelUIView.py +++ b/modules/ui/BaseConvertModelUIView.py @@ -1,56 +1,20 @@ -import traceback -from uuid import uuid4 - -from modules.util import create -from modules.util.args.ConvertModelArgs import ConvertModelArgs -from modules.util.config.TrainConfig import QuantizationConfig +from modules.util import path_util from modules.util.enum.DataType import DataType from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.ModelType import ModelType from modules.util.enum.PathIOType import PathIOType from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.ModelNames import EmbeddingName, ModelNames -from modules.util.torch_util import torch_gc -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState - -import customtkinter as ctk - - -class ConvertModelUI(ctk.CTkToplevel): - def __init__(self, parent, *args, **kwargs): - super().__init__(parent, *args, **kwargs) - self.parent = parent - - self.parent = parent - self.convert_model_args = ConvertModelArgs.default_values() - self.ui_state = UIState(self, self.convert_model_args) - self.button = None - - - self.title("Convert models") - self.geometry("550x350") - self.resizable(True, True) - self.frame = ctk.CTkFrame(self, width=600, height=300) - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) - self.main_frame(self.frame) - self.frame.pack(fill="both", expand=True) +class BaseConvertModelUIView: + def __init__(self, components): + self.components = components - self.wait_visibility() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def main_frame(self, master): + def build_content(self, frame, controller, ui_state): # model type - components.label(master, 0, 0, "Model Type", + self.components.label(frame, 0, 0, "Model Type", tooltip="Type of the model") - components.options_kv(master, 0, 1, [ #TODO simplify + self.components.options_kv(frame, 0, 1, [ #TODO simplify ("Stable Diffusion 1.5", ModelType.STABLE_DIFFUSION_15), ("Stable Diffusion 1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), ("Stable Diffusion 2.0", ModelType.STABLE_DIFFUSION_20), @@ -71,100 +35,49 @@ def main_frame(self, master): ("Chroma1", ModelType.CHROMA_1), #TODO does this just work? HiDream is not here ("QwenImage", ModelType.QWEN), #TODO does this just work? HiDream is not here ("ZImage", ModelType.Z_IMAGE), - ], self.ui_state, "model_type") + ], ui_state, "model_type") # training method - components.label(master, 1, 0, "Model Type", + self.components.label(frame, 1, 0, "Model Type", tooltip="The type of model to convert") - components.options_kv(master, 1, 1, [ + self.components.options_kv(frame, 1, 1, [ ("Base Model", TrainingMethod.FINE_TUNE), ("LoRA", TrainingMethod.LORA), ("Embedding", TrainingMethod.EMBEDDING), - ], self.ui_state, "training_method") + ], ui_state, "training_method") # input name - components.label(master, 2, 0, "Input name", + self.components.label(frame, 2, 0, "Input name", tooltip="Filename, directory or hugging face repository of the base model") - components.path_entry( - master, 2, 1, self.ui_state, "input_name", - mode="file", path_modifier=components.json_path_modifier + self.components.path_entry( + frame, 2, 1, ui_state, "input_name", + mode="file", path_modifier=path_util.json_path_modifier ) # output data type - components.label(master, 3, 0, "Output Data Type", + self.components.label(frame, 3, 0, "Output Data Type", tooltip="Precision to use when saving the output model") - components.options_kv(master, 3, 1, [ + self.components.options_kv(frame, 3, 1, [ ("float32", DataType.FLOAT_32), ("float16", DataType.FLOAT_16), ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "output_dtype") + ], ui_state, "output_dtype") # output format - components.label(master, 4, 0, "Output Format", + self.components.label(frame, 4, 0, "Output Format", tooltip="Format to use when saving the output model") - components.options_kv(master, 4, 1, [ + self.components.options_kv(frame, 4, 1, [ ("Safetensors", ModelFormat.SAFETENSORS), ("Diffusers", ModelFormat.DIFFUSERS), - ], self.ui_state, "output_model_format") + ], ui_state, "output_model_format") # output model destination - components.label(master, 5, 0, "Model Output Destination", + self.components.label(frame, 5, 0, "Model Output Destination", tooltip="Filename or directory where the output model is saved") - components.path_entry( - master, 5, 1, self.ui_state, "output_model_destination", + self.components.path_entry( + frame, 5, 1, ui_state, "output_model_destination", mode="file", io_type=PathIOType.MODEL, ) - self.button = components.button(master, 6, 1, "Convert", self.convert_model) - - def convert_model(self): - try: - self.button.configure(state="disabled") - model_loader = create.create_model_loader( - model_type=self.convert_model_args.model_type, - training_method=self.convert_model_args.training_method - ) - model_saver = create.create_model_saver( - model_type=self.convert_model_args.model_type, - training_method=self.convert_model_args.training_method - ) - - print("Loading model " + self.convert_model_args.input_name) - if self.convert_model_args.training_method in [TrainingMethod.FINE_TUNE]: - model = model_loader.load( - model_type=self.convert_model_args.model_type, - model_names=ModelNames( - base_model=self.convert_model_args.input_name, - ), - weight_dtypes=self.convert_model_args.weight_dtypes(), - quantization=QuantizationConfig.default_values(), - ) - elif self.convert_model_args.training_method in [TrainingMethod.LORA, TrainingMethod.EMBEDDING]: - model = model_loader.load( - model_type=self.convert_model_args.model_type, - model_names=ModelNames( - base_model=None, - lora=self.convert_model_args.input_name, - embedding=EmbeddingName(str(uuid4()), self.convert_model_args.input_name), - ), - weight_dtypes=self.convert_model_args.weight_dtypes(), - quantization=QuantizationConfig.default_values(), - ) - else: - raise Exception("could not load model: " + self.convert_model_args.input_name) - - print("Saving model " + self.convert_model_args.output_model_destination) - model_saver.save( - model=model, - model_type=self.convert_model_args.model_type, - output_model_format=self.convert_model_args.output_model_format, - output_model_destination=self.convert_model_args.output_model_destination, - dtype=self.convert_model_args.output_dtype.torch_dtype(), - ) - print("Model converted") - except Exception: - traceback.print_exc() - - torch_gc() - self.button.configure(state="normal") + self.button = self.components.button(frame, 6, 1, "Convert", controller.convert_model) diff --git a/modules/ui/BaseGenerateCaptionsWindowView.py b/modules/ui/BaseGenerateCaptionsWindowView.py index 1690879f1..68e24638e 100644 --- a/modules/ui/BaseGenerateCaptionsWindowView.py +++ b/modules/ui/BaseGenerateCaptionsWindowView.py @@ -1,133 +1,7 @@ -import contextlib -import tkinter as tk -from tkinter import filedialog +from abc import ABC, abstractmethod -from modules.util.ui.ui_utils import set_window_icon - -import customtkinter as ctk - - -class GenerateCaptionsWindow(ctk.CTkToplevel): - def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs): - """ - Window for generating captions for a folder of images - - Parameters: - parent (`Tk`): the parent window - path (`str`): the path to the folder - parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox - """ - super().__init__(parent, *args, **kwargs) - self.parent = parent - - if path is None: - path = "" - - self.mode_var = ctk.StringVar(self, "Create if absent") - self.modes = ["Replace all captions", "Create if absent", "Add as new line"] - self.model_var = ctk.StringVar(self, "Blip") - self.models = ["Blip", "Blip2", "WD14 VIT v2"] - - self.title("Batch generate captions") - self.geometry("360x360") - self.resizable(True, True) - - self.frame = ctk.CTkFrame(self, width=600, height=300) - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - - self.model_label = ctk.CTkLabel(self.frame, text="Model", width=100) - self.model_label.grid(row=0, column=0, sticky="w", padx=5, pady=5) - self.model_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.model_var, values=self.models, dynamic_resizing=False, width=200) - self.model_dropdown.grid(row=0, column=1, sticky="w", padx=5, pady=5) - - self.path_label = ctk.CTkLabel(self.frame, text="Folder", width=100) - self.path_label.grid(row=1, column=0, sticky="w",padx=5, pady=5) - self.path_entry = ctk.CTkEntry(self.frame, width=150) - self.path_entry.insert(0, path) - self.path_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.path_button = ctk.CTkButton(self.frame, width=30, text="...", command=lambda: self.browse_for_path(self.path_entry)) - self.path_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) - - self.caption_label = ctk.CTkLabel(self.frame, text="Initial Caption", width=100) - self.caption_label.grid(row=2, column=0, sticky="w", padx=5, pady=5) - self.caption_entry = ctk.CTkEntry(self.frame, width=200) - self.caption_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - - self.prefix_label = ctk.CTkLabel(self.frame, text="Caption Prefix", width=100) - self.prefix_label.grid(row=3, column=0, sticky="w", padx=5, pady=5) - self.prefix_entry = ctk.CTkEntry(self.frame, width=200) - self.prefix_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) - - self.postfix_label = ctk.CTkLabel(self.frame, text="Caption Postfix", width=100) - self.postfix_label.grid(row=4, column=0, sticky="w", padx=5, pady=5) - self.postfix_entry = ctk.CTkEntry(self.frame, width=200) - self.postfix_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) - - self.mode_label = ctk.CTkLabel(self.frame, text="Mode", width=100) - self.mode_label.grid(row=5, column=0, sticky="w", padx=5, pady=5) - self.mode_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.mode_var, values=self.modes, dynamic_resizing=False, width=200) - self.mode_dropdown.grid(row=5, column=1, sticky="w", padx=5, pady=5) - - self.include_subdirectories_label = ctk.CTkLabel(self.frame, text="Include subfolders", width=100) - self.include_subdirectories_label.grid(row=6, column=0, sticky="w", padx=5, pady=5) - self.include_subdirectories_var = ctk.BooleanVar(self, parent_include_subdirectories) - self.include_subdirectories_switch = ctk.CTkSwitch(self.frame, text="", variable=self.include_subdirectories_var) - self.include_subdirectories_switch.grid(row=6, column=1, sticky="w", padx=5, pady=5) - - self.progress_label = ctk.CTkLabel(self.frame, text="Progress: 0/0", width=100) - self.progress_label.grid(row=7, column=0, sticky="w", padx=5, pady=5) - self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) - self.progress.grid(row=7, column=1, sticky="w", padx=5, pady=5) - - self.create_captions_button = ctk.CTkButton(self.frame, text="Create Captions", width=310, command=self.create_captions) - self.create_captions_button.grid(row=8, column=0, columnspan=2, sticky="w", padx=5, pady=5) - - self.frame.pack(fill="both", expand=True) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def browse_for_path(self, entry_box): - # get the path from the user - path = filedialog.askdirectory() - # set the path to the entry box - # delete entry box text - entry_box.focus_set() - entry_box.delete(0, filedialog.END) - entry_box.insert(0, path) - self.focus_set() +class BaseGenerateCaptionsWindowView(ABC): + @abstractmethod def set_progress(self, value, max_value): - progress = value / max_value - self.progress.set(progress) - self.progress_label.configure(text=f"{value}/{max_value}") - self.progress.update() - - def create_captions(self): - self.parent.load_captioning_model(self.model_var.get()) - - mode = { - "Replace all captions": "replace", - "Create if absent": "fill", - "Add as new line": "add", - }[self.mode_var.get()] - - self.parent.captioning_model.caption_folder( - sample_dir=self.path_entry.get(), - initial_caption=self.caption_entry.get(), - caption_prefix=self.prefix_entry.get(), - caption_postfix=self.postfix_entry.get(), - mode=mode, - progress_callback=self.set_progress, - include_subdirectories=self.include_subdirectories_var.get(), - ) - self.parent.load_image() - - def destroy(self): - with contextlib.suppress(tk.TclError): - self.grab_release() - - super().destroy() + pass diff --git a/modules/ui/BaseGenerateMasksWindowView.py b/modules/ui/BaseGenerateMasksWindowView.py index daff0d3d5..f9e82231e 100644 --- a/modules/ui/BaseGenerateMasksWindowView.py +++ b/modules/ui/BaseGenerateMasksWindowView.py @@ -1,151 +1,7 @@ -import contextlib -import tkinter as tk -from tkinter import filedialog +from abc import ABC, abstractmethod -from modules.util.ui.ui_utils import set_window_icon - -import customtkinter as ctk - - -class GenerateMasksWindow(ctk.CTkToplevel): - def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs): - """ - Window for generating masks for a folder of images - - Parameters: - parent (`Tk`): the parent window - path (`str`): the path to the folder - parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox - """ - super().__init__(parent, *args, **kwargs) - - self.parent = parent - if path is None: - path = "" - - self.mode_var = ctk.StringVar(self, "Create if absent") - self.modes = ["Replace all masks", "Create if absent", "Add to existing", "Subtract from existing", "Blend with existing"] - self.model_var = ctk.StringVar(self, "ClipSeg") - self.models = ["ClipSeg", "Rembg", "Rembg-Human", "Hex Color"] - - self.title("Batch generate masks") - self.geometry("360x430") - self.resizable(True, True) - - self.frame = ctk.CTkFrame(self, width=600, height=300) - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - - self.model_label = ctk.CTkLabel(self.frame, text="Model", width=100) - self.model_label.grid(row=0, column=0, sticky="w", padx=5, pady=5) - self.model_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.model_var, values=self.models, dynamic_resizing=False, width=200) - self.model_dropdown.grid(row=0, column=1, sticky="w", padx=5, pady=5) - - self.path_label = ctk.CTkLabel(self.frame, text="Folder", width=100) - self.path_label.grid(row=1, column=0, sticky="w",padx=5, pady=5) - self.path_entry = ctk.CTkEntry(self.frame, width=150) - self.path_entry.insert(0, path) - self.path_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.path_button = ctk.CTkButton(self.frame, width=30, text="...", command=lambda: self.browse_for_path(self.path_entry)) - self.path_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) - - self.prompt_label = ctk.CTkLabel(self.frame, text="Prompt", width=100) - self.prompt_label.grid(row=2, column=0, sticky="w",padx=5, pady=5) - self.prompt_entry = ctk.CTkEntry(self.frame, width=200) - self.prompt_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - - self.mode_label = ctk.CTkLabel(self.frame, text="Mode", width=100) - self.mode_label.grid(row=3, column=0, sticky="w", padx=5, pady=5) - self.mode_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.mode_var, values=self.modes, dynamic_resizing=False, width=200) - self.mode_dropdown.grid(row=3, column=1, sticky="w", padx=5, pady=5) - - self.threshold_label = ctk.CTkLabel(self.frame, text="Threshold", width=100) - self.threshold_label.grid(row=4, column=0, sticky="w", padx=5, pady=5) - self.threshold_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="0.0 - 1.0") - self.threshold_entry.insert(0, "0.3") - self.threshold_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) - - self.smooth_label = ctk.CTkLabel(self.frame, text="Smooth", width=100) - self.smooth_label.grid(row=5, column=0, sticky="w", padx=5, pady=5) - self.smooth_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="5") - self.smooth_entry.insert(0, 5) - self.smooth_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) - - self.expand_label = ctk.CTkLabel(self.frame, text="Expand", width=100) - self.expand_label.grid(row=6, column=0, sticky="w", padx=5, pady=5) - self.expand_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="10") - self.expand_entry.insert(0, 10) - self.expand_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) - - self.alpha_label = ctk.CTkLabel(self.frame, text="Alpha", width=100) - self.alpha_label.grid(row=7, column=0, sticky="w", padx=5, pady=5) - self.alpha_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="1") - self.alpha_entry.insert(0, 1) - self.alpha_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) - - self.include_subdirectories_label = ctk.CTkLabel(self.frame, text="Include subfolders", width=100) - self.include_subdirectories_label.grid(row=8, column=0, sticky="w", padx=5, pady=5) - self.include_subdirectories_var = ctk.BooleanVar(self, parent_include_subdirectories) - self.include_subdirectories_switch = ctk.CTkSwitch(self.frame, text="", variable=self.include_subdirectories_var) - self.include_subdirectories_switch.grid(row=8, column=1, sticky="w", padx=5, pady=5) - - self.progress_label = ctk.CTkLabel(self.frame, text="Progress: 0/0", width=100) - self.progress_label.grid(row=9, column=0, sticky="w", padx=5, pady=5) - self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) - self.progress.grid(row=9, column=1, sticky="w", padx=5, pady=5) - - self.create_masks_button = ctk.CTkButton(self.frame, text="Create Masks", width=310, command=self.create_masks) - self.create_masks_button.grid(row=10, column=0, columnspan=2, sticky="w", padx=5, pady=5) - - self.frame.pack(fill="both", expand=True) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def browse_for_path(self, entry_box): - # get the path from the user - path = filedialog.askdirectory() - # set the path to the entry box - # delete entry box text - entry_box.focus_set() - entry_box.delete(0, filedialog.END) - entry_box.insert(0, path) - self.focus_set() +class BaseGenerateMasksWindowView(ABC): + @abstractmethod def set_progress(self, value, max_value): - progress = value / max_value - self.progress.set(progress) - self.progress_label.configure(text=f"{value}/{max_value}") - self.progress.update() - - def create_masks(self): - self.parent.load_masking_model(self.model_var.get()) - - mode = { - "Replace all masks": "replace", - "Create if absent": "fill", - "Add to existing": "add", - "Subtract from existing": "subtract", - "Blend with existing": "blend", - }[self.mode_var.get()] - - self.parent.masking_model.mask_folder( - sample_dir=self.path_entry.get(), - prompts=[self.prompt_entry.get()], - mode=mode, - alpha=float(self.alpha_entry.get()), - threshold=float(self.threshold_entry.get()), - smooth_pixels=int(self.smooth_entry.get()), - expand_pixels=int(self.expand_entry.get()), - progress_callback=self.set_progress, - include_subdirectories=self.include_subdirectories_var.get(), - ) - self.parent.load_image() - - def destroy(self): - with contextlib.suppress(tk.TclError): - self.grab_release() - - super().destroy() + pass diff --git a/modules/ui/BaseLoraTabView.py b/modules/ui/BaseLoraTabView.py index 1c73d90ce..ea6f52d5c 100644 --- a/modules/ui/BaseLoraTabView.py +++ b/modules/ui/BaseLoraTabView.py @@ -1,47 +1,20 @@ -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.DataType import DataType +from modules.util import path_util from modules.util.enum.ModelType import PeftType -from modules.util.ui import components -from modules.util.ui.UIState import UIState from modules.util.ui.validation_helpers import check_range -import customtkinter as ctk +class BaseLoraTabView: + def __init__(self, components): + self.components = components -class LoraTab: - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__() + def build(self, frame, controller, ui_state, setup_lora_callback): + self.components.label(frame, 0, 0, "Type", + tooltip="The type of low-parameter finetuning method.") + self.components.options_kv(frame, 0, 1, controller.get_peft_types(), + ui_state, "peft_type", command=setup_lora_callback) - self.master = master - self.train_config = train_config - self.ui_state = ui_state - - self.scroll_frame = None - self.options_frame = None - - self.refresh_ui() - - def refresh_ui(self): - if self.scroll_frame: - self.scroll_frame.destroy() - self.scroll_frame = ctk.CTkFrame(self.master, fg_color="transparent") - self.scroll_frame.grid(row=0, column=0, sticky="nsew") - - self.scroll_frame.grid_columnconfigure(0, weight=0) - self.scroll_frame.grid_columnconfigure(1, weight=1) - self.scroll_frame.grid_columnconfigure(2, weight=2) - - components.label(self.scroll_frame, 0, 0, "Type", - tooltip="The type of low-parameter finetuning method.") - # This will instantly call self.setup_lora. - components.options_kv(self.scroll_frame, 0, 1, [ - ("LoRA", PeftType.LORA), - ("LoHa", PeftType.LOHA), - ("OFT v2", PeftType.OFT_2), - ], self.ui_state, "peft_type", command=self.setup_lora) - - def setup_lora(self, peft_type: PeftType): + def build_lora_options(self, master, controller, ui_state, peft_type: PeftType): if peft_type == PeftType.LOHA: name = "LoHa" elif peft_type == PeftType.OFT_2: @@ -49,106 +22,87 @@ def setup_lora(self, peft_type: PeftType): else: name = "LoRA" - if self.options_frame: - self.options_frame.destroy() - self.options_frame = ctk.CTkFrame(self.scroll_frame, fg_color="transparent") - self.options_frame.grid(row=1, column=0, columnspan=3, sticky="nsew") - master = self.options_frame - - master.grid_columnconfigure(0, weight=0, uniform="a") - master.grid_columnconfigure(1, weight=1, uniform="a") - master.grid_columnconfigure(2, minsize=50, uniform="a") - master.grid_columnconfigure(3, weight=0, uniform="a") - master.grid_columnconfigure(4, weight=1, uniform="a") - # 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.path_entry( - master, 0, 1, self.ui_state, "lora_model_name", - mode="file", path_modifier=components.json_path_modifier + self.components.label(master, 0, 0, f"{name} base model", + tooltip=f"The base {name} to train on. Leave empty to create a new {name}") + self.components.path_entry( + master, 0, 1, ui_state, "lora_model_name", + mode="file", path_modifier=path_util.json_path_modifier, + columnspan=4, ) - entry.grid(row=0, column=1, columnspan=4) - # LoRA decomposition if peft_type == PeftType.LORA: - components.label(master, 1, 3, "Decompose Weights (DoRA)", - tooltip="Decompose LoRA Weights (aka, DoRA).") - components.switch(master, 1, 4, self.ui_state, "lora_decompose") + self.components.label(master, 1, 3, "Decompose Weights (DoRA)", + tooltip="Decompose LoRA Weights (aka, DoRA).") + self.components.switch(master, 1, 4, ui_state, "lora_decompose") - components.label(master, 2, 3, "Use Norm Epsilon (DoRA Only)", - tooltip="Add an epsilon to the norm divison calculation in DoRA. Can aid in training stability, and also acts as regularization.") - components.switch(master, 2, 4, self.ui_state, "lora_decompose_norm_epsilon") - components.label(master, 3, 3, "Apply on output axis (DoRA Only)", - tooltip="Apply the weight decomposition on the output axis instead of the input axis.") - components.switch(master, 3, 4, self.ui_state, "lora_decompose_output_axis") + self.components.label(master, 2, 3, "Use Norm Epsilon (DoRA Only)", + tooltip="Add an epsilon to the norm divison calculation in DoRA. Can aid in training stability, and also acts as regularization.") + self.components.switch(master, 2, 4, ui_state, "lora_decompose_norm_epsilon") + self.components.label(master, 3, 3, "Apply on output axis (DoRA Only)", + tooltip="Apply the weight decomposition on the output axis instead of the input axis.") + self.components.switch(master, 3, 4, ui_state, "lora_decompose_output_axis") # LoRA and LoHA shared settings if peft_type == PeftType.LORA or peft_type == PeftType.LOHA: # rank - components.label(master, 1, 0, f"{name} rank", - tooltip=f"The rank parameter used when creating a new {name}") - components.entry(master, 1, 1, self.ui_state, "lora_rank", required=True, extra_validate=check_range(lower=1, message="Rank must be at least 1")) + self.components.label(master, 1, 0, f"{name} rank", + tooltip=f"The rank parameter used when creating a new {name}") + self.components.entry(master, 1, 1, ui_state, "lora_rank", required=True, extra_validate=check_range(lower=1, message="Rank must be at least 1")) # alpha - components.label(master, 2, 0, f"{name} alpha", - tooltip=f"The alpha parameter used when creating a new {name}") - components.entry(master, 2, 1, self.ui_state, "lora_alpha", required=True) + self.components.label(master, 2, 0, f"{name} alpha", + tooltip=f"The alpha parameter used when creating a new {name}") + self.components.entry(master, 2, 1, ui_state, "lora_alpha", required=True) # Dropout Percentage - components.label(master, 3, 0, "Dropout Probability", - tooltip="Dropout probability. This percentage of model nodes will be randomly ignored at each training step. Helps with overfitting. 0 disables, 1 maximum.") - components.entry(master, 3, 1, self.ui_state, "dropout_probability") + self.components.label(master, 3, 0, "Dropout Probability", + tooltip="Dropout probability. This percentage of model nodes will be randomly ignored at each training step. Helps with overfitting. 0 disables, 1 maximum.") + self.components.entry(master, 3, 1, ui_state, "dropout_probability") # weight dtype - components.label(master, 4, 0, f"{name} Weight Data Type", - tooltip=f"The {name} weight data type used for training. This can reduce memory consumption, but reduces precision") - components.options_kv(master, 4, 1, [ - ("float32", DataType.FLOAT_32), - ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "lora_weight_dtype") + self.components.label(master, 4, 0, f"{name} Weight Data Type", + tooltip=f"The {name} weight data type used for training. This can reduce memory consumption, but reduces precision") + self.components.options_kv(master, 4, 1, controller.get_lora_weight_dtypes(), ui_state, "lora_weight_dtype") # For use with additional embeddings. - components.label(master, 5, 0, "Bundle Embeddings", - tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") - components.switch(master, 5, 1, self.ui_state, "bundle_additional_embeddings") + self.components.label(master, 5, 0, "Bundle Embeddings", + tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") + self.components.switch(master, 5, 1, ui_state, "bundle_additional_embeddings") # OFTv2 elif peft_type == PeftType.OFT_2: # Block Size - components.label(master, 1, 0, f"{name} Block Size", - tooltip=f"The block size parameter used when creating a new {name}") - components.entry(master, 1, 1, self.ui_state, "oft_block_size", required=True) + self.components.label(master, 1, 0, f"{name} Block Size", + tooltip=f"The block size parameter used when creating a new {name}") + self.components.entry(master, 1, 1, ui_state, "oft_block_size", required=True) # COFT - components.label(master, 1, 3, "Constrained OFT (COFT)", - tooltip="Use the constrained variant of OFT. This constrains the learned rotation to stay very close to the identity matrix, limiting adaptation to only small changes. This improves training stability, helps prevent overfitting on small datasets, and better preserves the base models original knowledge but it may lack expressiveness for tasks requiring substantial adaptation and introduces an additional hyperparameter (COFT Epsilon) that needs tuning.") - components.switch(master, 1, 4, self.ui_state, "oft_coft") + self.components.label(master, 1, 3, "Constrained OFT (COFT)", + tooltip="Use the constrained variant of OFT. This constrains the learned rotation to stay very close to the identity matrix, limiting adaptation to only small changes. This improves training stability, helps prevent overfitting on small datasets, and better preserves the base models original knowledge but it may lack expressiveness for tasks requiring substantial adaptation and introduces an additional hyperparameter (COFT Epsilon) that needs tuning.") + self.components.switch(master, 1, 4, ui_state, "oft_coft") - components.label(master, 2, 3, "COFT Epsilon", - tooltip="The control strength of COFT. Only has an effect if COFT is enabled.") - components.entry(master, 2, 4, self.ui_state, "coft_eps") + self.components.label(master, 2, 3, "COFT Epsilon", + tooltip="The control strength of COFT. Only has an effect if COFT is enabled.") + self.components.entry(master, 2, 4, ui_state, "coft_eps") # Block Share - components.label(master, 3, 3, "Block Share", - tooltip="Share the OFT parameters between blocks. A single rotation matrix is shared across all blocks within a layer, drastically cutting the number of trainable parameters and yielding very compact adapter files, potentially improving generalization but at the cost of significant expressiveness, which can lead to underfitting on more complex or diverse tasks.") - components.switch(master, 3, 4, self.ui_state, "oft_block_share") + self.components.label(master, 3, 3, "Block Share", + tooltip="Share the OFT parameters between blocks. A single rotation matrix is shared across all blocks within a layer, drastically cutting the number of trainable parameters and yielding very compact adapter files, potentially improving generalization but at the cost of significant expressiveness, which can lead to underfitting on more complex or diverse tasks.") + self.components.switch(master, 3, 4, ui_state, "oft_block_share") # Dropout Percentage - components.label(master, 2, 0, "Dropout Probability", - tooltip="Dropout probability. This percentage of the rotated adapter nodes that will be randomly restored to the base model initial statue. Helps with overfitting. 0 disables, 1 maximum.") - components.entry(master, 2, 1, self.ui_state, "dropout_probability") + self.components.label(master, 2, 0, "Dropout Probability", + tooltip="Dropout probability. This percentage of the rotated adapter nodes that will be randomly restored to the base model initial statue. Helps with overfitting. 0 disables, 1 maximum.") + self.components.entry(master, 2, 1, ui_state, "dropout_probability") # OFT weight dtype - components.label(master, 3, 0, f"{name} Weight Data Type", - tooltip=f"The {name} weight data type used for training. This can reduce memory consumption, but reduces precision") - components.options_kv(master, 3, 1, [ - ("float32", DataType.FLOAT_32), - ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "lora_weight_dtype") + self.components.label(master, 3, 0, f"{name} Weight Data Type", + tooltip=f"The {name} weight data type used for training. This can reduce memory consumption, but reduces precision") + self.components.options_kv(master, 3, 1, controller.get_lora_weight_dtypes(), ui_state, "lora_weight_dtype") # For use with additional embeddings. - components.label(master, 4, 0, "Bundle Embeddings", - tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") - components.switch(master, 4, 1, self.ui_state, "bundle_additional_embeddings") + self.components.label(master, 4, 0, "Bundle Embeddings", + tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") + self.components.switch(master, 4, 1, ui_state, "bundle_additional_embeddings") diff --git a/modules/ui/BaseModelTabView.py b/modules/ui/BaseModelTabView.py index ff17ea3ba..88f223e22 100644 --- a/modules/ui/BaseModelTabView.py +++ b/modules/ui/BaseModelTabView.py @@ -1,85 +1,60 @@ -from modules.util import create -from modules.util.config.TrainConfig import TrainConfig +from abc import ABC, abstractmethod + +from modules.util import path_util from modules.util.enum.ConfigPart import ConfigPart from modules.util.enum.DataType import DataType from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.PathIOType import PathIOType from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.ui import components -from modules.util.ui.UIState import UIState - -import customtkinter as ctk - - -class ModelTab: - - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__() - - self.master = master - self.train_config = train_config - self.ui_state = ui_state - - master.grid_rowconfigure(0, weight=1) - master.grid_columnconfigure(0, weight=1) - - self.scroll_frame = None - - self.refresh_ui() - - def refresh_ui(self): - if self.scroll_frame: - self.scroll_frame.destroy() - - self.scroll_frame = ctk.CTkScrollableFrame(self.master, fg_color="transparent") - self.scroll_frame.grid(row=0, column=0, sticky="nsew") - self.scroll_frame.grid_columnconfigure(0, weight=1) - - base_frame = ctk.CTkFrame(master=self.scroll_frame, corner_radius=5) - base_frame.grid(row=0, column=0, padx=5, pady=5, sticky="nsew") - - base_frame.grid_columnconfigure(0, weight=0) - base_frame.grid_columnconfigure(1, weight=10)#, minsize=500) - base_frame.grid_columnconfigure(2, minsize=50) - base_frame.grid_columnconfigure(3, weight=0) - base_frame.grid_columnconfigure(4, weight=1) - - if self.train_config.model_type.is_stable_diffusion(): #TODO simplify - self.__setup_stable_diffusion_ui(base_frame) - if self.train_config.model_type.is_stable_diffusion_3(): - self.__setup_stable_diffusion_3_ui(base_frame) - elif self.train_config.model_type.is_stable_diffusion_xl(): - self.__setup_stable_diffusion_xl_ui(base_frame) - elif self.train_config.model_type.is_wuerstchen(): - self.__setup_wuerstchen_ui(base_frame) - elif self.train_config.model_type.is_pixart(): - self.__setup_pixart_alpha_ui(base_frame) - elif self.train_config.model_type.is_flux_1(): - self.__setup_flux_ui(base_frame) - elif self.train_config.model_type.is_flux_2(): - self.__setup_flux_2_ui(base_frame) - elif self.train_config.model_type.is_z_image(): - self.__setup_z_image_ui(base_frame) - elif self.train_config.model_type.is_chroma(): - self.__setup_chroma_ui(base_frame) - elif self.train_config.model_type.is_qwen(): - self.__setup_qwen_ui(base_frame) - elif self.train_config.model_type.is_sana(): - self.__setup_sana_ui(base_frame) - elif self.train_config.model_type.is_hunyuan_video(): - self.__setup_hunyuan_video_ui(base_frame) - elif self.train_config.model_type.is_hi_dream(): - self.__setup_hi_dream_ui(base_frame) - elif self.train_config.model_type.is_ernie(): - self.__setup_ernie_ui(base_frame) - - def __setup_stable_diffusion_ui(self, frame): + + +class BaseModelTabView(ABC): + def __init__(self, components): + self.components = components + + @abstractmethod + def _make_svd_frames(self, parent, row: int): + """Create and place SVDQuant label+entry subframes; return (label_frame, entry_frame).""" + + def build_content(self, frame, controller, ui_state): + if controller.train_config.model_type.is_stable_diffusion(): # TODO simplify + self.__setup_stable_diffusion_ui(frame, controller, ui_state) + if controller.train_config.model_type.is_stable_diffusion_3(): + self.__setup_stable_diffusion_3_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_stable_diffusion_xl(): + self.__setup_stable_diffusion_xl_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_wuerstchen(): + self.__setup_wuerstchen_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_pixart(): + self.__setup_pixart_alpha_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_flux_1(): + self.__setup_flux_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_flux_2(): + self.__setup_flux_2_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_z_image(): + self.__setup_z_image_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_chroma(): + self.__setup_chroma_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_qwen(): + self.__setup_qwen_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_sana(): + self.__setup_sana_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_hunyuan_video(): + self.__setup_hunyuan_video_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_hi_dream(): + self.__setup_hi_dream_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_ernie(): + self.__setup_ernie_ui(frame, controller, ui_state) + + def __setup_stable_diffusion_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_unet=True, has_text_encoder=True, has_vae=True, @@ -87,20 +62,23 @@ def __setup_stable_diffusion_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method in [ + allow_diffusers=controller.train_config.training_method in [ TrainingMethod.FINE_TUNE, TrainingMethod.FINE_TUNE_VAE, ], - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_stable_diffusion_3_ui(self, frame): + def __setup_stable_diffusion_3_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, has_text_encoder_1=True, has_text_encoder_2=True, @@ -110,17 +88,20 @@ def __setup_stable_diffusion_3_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_flux_ui(self, frame): + def __setup_flux_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, allow_override_transformer=True, has_text_encoder_1=True, @@ -130,17 +111,20 @@ def __setup_flux_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_flux_2_ui(self, frame): + def __setup_flux_2_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, allow_override_transformer=True, has_text_encoder_1=True, @@ -149,17 +133,20 @@ def __setup_flux_2_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_z_image_ui(self, frame): + def __setup_z_image_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, allow_override_transformer=True, has_text_encoder_1=True, @@ -168,17 +155,20 @@ def __setup_z_image_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_ernie_ui(self, frame): + def __setup_ernie_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, allow_override_transformer=True, has_text_encoder_1=True, @@ -187,17 +177,20 @@ def __setup_ernie_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_chroma_ui(self, frame): + def __setup_chroma_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, allow_override_transformer=True, has_text_encoder_1=True, @@ -206,17 +199,20 @@ def __setup_chroma_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_qwen_ui(self, frame): + def __setup_qwen_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, allow_override_transformer=True, has_text_encoder_1=True, @@ -225,17 +221,20 @@ def __setup_qwen_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_stable_diffusion_xl_ui(self, frame): + def __setup_stable_diffusion_xl_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_unet=True, has_text_encoder_1=True, has_text_encoder_2=True, @@ -244,38 +243,44 @@ def __setup_stable_diffusion_xl_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_wuerstchen_ui(self, frame): + def __setup_wuerstchen_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_prior=True, - allow_override_prior=self.train_config.model_type.is_stable_cascade(), + allow_override_prior=controller.train_config.model_type.is_stable_cascade(), has_text_encoder=True, ) - row = self.__create_effnet_encoder_components(frame, row) - row = self.__create_decoder_components(frame, row, self.train_config.model_type.is_wuerstchen_v2()) + row = self.__create_effnet_encoder_components(frame, row, ui_state) + row = self.__create_decoder_components(frame, row, ui_state, controller.train_config.model_type.is_wuerstchen_v2()) row = self.__create_output_components( frame, row, - allow_safetensors=self.train_config.training_method != TrainingMethod.FINE_TUNE - or self.train_config.model_type.is_stable_cascade(), - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ui_state, + allow_safetensors=controller.train_config.training_method != TrainingMethod.FINE_TUNE + or controller.train_config.model_type.is_stable_cascade(), + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_pixart_alpha_ui(self, frame): + def __setup_pixart_alpha_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, has_text_encoder=True, has_vae=True, @@ -283,17 +288,20 @@ def __setup_pixart_alpha_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_sana_ui(self, frame): + def __setup_sana_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, has_text_encoder=True, has_vae=True, @@ -301,17 +309,20 @@ def __setup_sana_ui(self, frame): row = self.__create_output_components( frame, row, - allow_safetensors=self.train_config.training_method != TrainingMethod.FINE_TUNE, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ui_state, + allow_safetensors=controller.train_config.training_method != TrainingMethod.FINE_TUNE, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_hunyuan_video_ui(self, frame): + def __setup_hunyuan_video_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, allow_override_transformer=True, has_text_encoder_1=True, @@ -321,17 +332,20 @@ def __setup_hunyuan_video_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_hi_dream_ui(self, frame): + def __setup_hi_dream_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, has_text_encoder_1=True, has_text_encoder_2=True, @@ -343,12 +357,13 @@ def __setup_hi_dream_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __create_dtype_options(self, include_gguf: bool=False, include_a8: bool=False) -> list[tuple[str, DataType]]: + def __create_dtype_options(self, include_gguf: bool = False, include_a8: bool = False) -> list[tuple[str, DataType]]: options = [ ("float32", DataType.FLOAT_32), ("bfloat16", DataType.BFLOAT_16), @@ -373,28 +388,28 @@ def __create_dtype_options(self, include_gguf: bool=False, include_a8: bool=Fals return options - def __create_base_dtype_components(self, frame, row: int) -> int: + def __create_base_dtype_components(self, frame, row: int, ui_state) -> int: # huggingface token - components.label(frame, row, 0, "Hugging Face Token", + self.components.label(frame, row, 0, "Hugging Face Token", tooltip="Enter your Hugging Face access token if you have used a protected Hugging Face repository below.\nThis value is stored separately, not saved to your configuration file. " "Go to https://huggingface.co/settings/tokens to create an access token.", wide_tooltip=True) - components.entry(frame, row, 1, self.ui_state, "secrets.huggingface_token") + self.components.entry(frame, row, 1, ui_state, "secrets.huggingface_token") row += 1 # base model - components.label(frame, row, 0, "Base Model", + self.components.label(frame, row, 0, "Base Model", tooltip="Filename, directory or Hugging Face repository of the base model") - components.path_entry( - frame, row, 1, self.ui_state, "base_model_name", - mode="file", path_modifier=components.json_path_modifier + self.components.path_entry( + frame, row, 1, ui_state, "base_model_name", + mode="file", path_modifier=path_util.json_path_modifier ) # compile - components.label(frame, row, 3, "Compile transformer blocks", + self.components.label(frame, row, 3, "Compile transformer blocks", tooltip="Uses torch.compile and Triton to significantly speed up training. Only applies to transformer/unet. Disable in case of compatibility issues.") - components.switch(frame, row, 4, self.ui_state, "compile") + self.components.switch(frame, row, 4, ui_state, "compile") row += 1 @@ -404,6 +419,8 @@ def __create_base_components( self, frame, row: int, + controller, + ui_state, has_unet: bool = False, has_prior: bool = False, allow_override_prior: bool = False, @@ -419,54 +436,53 @@ def __create_base_components( ) -> int: if has_unet: # unet weight dtype - components.label(frame, row, 3, "UNet Data Type", + self.components.label(frame, row, 3, "UNet Data Type", tooltip="The unet weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(include_a8=True), - self.ui_state, "unet.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(include_a8=True), + ui_state, "unet.weight_dtype") row += 1 if has_prior: if allow_override_prior: # prior model - components.label(frame, row, 0, "Prior Model", + self.components.label(frame, row, 0, "Prior Model", tooltip="Filename, directory or Hugging Face repository of the prior model") - components.path_entry( - frame, row, 1, self.ui_state, "prior.model_name", - mode="file", path_modifier=components.json_path_modifier + self.components.path_entry( + frame, row, 1, ui_state, "prior.model_name", + mode="file", path_modifier=path_util.json_path_modifier ) # prior weight dtype - components.label(frame, row, 3, "Prior Data Type", + self.components.label(frame, row, 3, "Prior Data Type", tooltip="The prior weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "prior.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "prior.weight_dtype") row += 1 if has_transformer: if allow_override_transformer: # transformer model - components.label(frame, row, 0, "Override Transformer / GGUF", + self.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.path_entry( - frame, row, 1, self.ui_state, "transformer.model_name", - mode="file", path_modifier=components.json_path_modifier + self.components.path_entry( + frame, row, 1, ui_state, "transformer.model_name", + mode="file", path_modifier=path_util.json_path_modifier ) # transformer weight dtype - components.label(frame, row, 3, "Transformer Data Type", + self.components.label(frame, row, 3, "Transformer Data Type", tooltip="The transformer weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(include_gguf=True, include_a8=True), - self.ui_state, "transformer.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(include_gguf=True, include_a8=True), + ui_state, "transformer.weight_dtype") row += 1 - cls = create.get_model_setup_class(self.train_config.model_type, self.train_config.training_method) - presets = cls.LAYER_PRESETS if cls is not None else {"full": []} + presets = controller.get_presets() - components.label(frame, row, 0, "Quantization") - components.layer_filter_entry(frame, row, 1, self.ui_state, + self.components.label(frame, row, 0, "Quantization") + self.components.layer_filter_entry(frame, row, 1, ui_state, preset_var_name="quantization.layer_filter_preset", presets=presets, preset_label="Quantization Layer Filter", preset_tooltip="Select a preset defining which layers to quantize. Quantization of certain layers can decrease model quality. Only applies to the transformer/unet", @@ -478,107 +494,103 @@ def __create_base_components( ) # SVDQuant - create vertical grids to match the size of layer_filter_entry - svd_label_frame = ctk.CTkFrame(frame, fg_color="transparent") - svd_label_frame.grid(row=row, column=3, sticky="nsew") - svd_entry_frame = ctk.CTkFrame(frame, fg_color="transparent") - svd_entry_frame.grid(row=row, column=4, sticky="nsew") - components.label(svd_label_frame, 0, 0, "SVDQuant", + svd_label_frame, svd_entry_frame = self._make_svd_frames(frame, row) + self.components.label(svd_label_frame, 0, 0, "SVDQuant", tooltip="What datatype to use for SVDQuant weights decomposition.") - components.options_kv(svd_entry_frame, 0, 0, [("disabled", DataType.NONE), ("float32", DataType.FLOAT_32), ("bfloat16", DataType.BFLOAT_16)], - self.ui_state, "quantization.svd_dtype") - components.label(svd_label_frame, 1, 0, "SVDQuant Rank", + self.components.options_kv(svd_entry_frame, 0, 0, [("disabled", DataType.NONE), ("float32", DataType.FLOAT_32), ("bfloat16", DataType.BFLOAT_16)], + ui_state, "quantization.svd_dtype") + self.components.label(svd_label_frame, 1, 0, "SVDQuant Rank", tooltip="Rank for SVDQuant weights decomposition") - components.entry(svd_entry_frame, 1, 0, self.ui_state, "quantization.svd_rank") + self.components.entry(svd_entry_frame, 1, 0, ui_state, "quantization.svd_rank") row += 1 - if has_text_encoder: # text encoder weight dtype - components.label(frame, row, 3, "Text Encoder Data Type", + self.components.label(frame, row, 3, "Text Encoder Data Type", tooltip="The text encoder weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "text_encoder.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "text_encoder.weight_dtype") row += 1 if has_text_encoder_1: # text encoder 1 weight dtype - components.label(frame, row, 3, "Text Encoder 1 Data Type", + self.components.label(frame, row, 3, "Text Encoder 1 Data Type", tooltip="The text encoder 1 weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "text_encoder.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "text_encoder.weight_dtype") row += 1 if has_text_encoder_2: # text encoder 2 weight dtype - components.label(frame, row, 3, "Text Encoder 2 Data Type", + self.components.label(frame, row, 3, "Text Encoder 2 Data Type", tooltip="The text encoder 2 weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "text_encoder_2.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "text_encoder_2.weight_dtype") row += 1 if has_text_encoder_3: # text encoder 3 weight dtype - components.label(frame, row, 3, "Text Encoder 3 Data Type", + self.components.label(frame, row, 3, "Text Encoder 3 Data Type", tooltip="The text encoder 3 weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "text_encoder_3.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "text_encoder_3.weight_dtype") row += 1 if has_text_encoder_4: if allow_override_text_encoder_4: # text encoder 4 weight dtype - components.label(frame, row, 0, "Text Encoder 4 Override", + self.components.label(frame, row, 0, "Text Encoder 4 Override", tooltip="Filename, directory or Hugging Face repository of the text encoder 4 model") - components.path_entry( - frame, row, 1, self.ui_state, "text_encoder_4.model_name", - mode="file", path_modifier=components.json_path_modifier + self.components.path_entry( + frame, row, 1, ui_state, "text_encoder_4.model_name", + mode="file", path_modifier=path_util.json_path_modifier ) # text encoder 4 weight dtype - components.label(frame, row, 3, "Text Encoder 4 Data Type", + self.components.label(frame, row, 3, "Text Encoder 4 Data Type", tooltip="The text encoder 4 weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "text_encoder_4.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "text_encoder_4.weight_dtype") row += 1 if has_vae: # base model - components.label(frame, row, 0, "VAE Override", + self.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.path_entry( - frame, row, 1, self.ui_state, "vae.model_name", - mode="file", path_modifier=components.json_path_modifier + self.components.path_entry( + frame, row, 1, ui_state, "vae.model_name", + mode="file", path_modifier=path_util.json_path_modifier ) # vae weight dtype - components.label(frame, row, 3, "VAE Data Type", + self.components.label(frame, row, 3, "VAE Data Type", tooltip="The vae weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "vae.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "vae.weight_dtype") row += 1 return row - def __create_effnet_encoder_components(self, frame, row: int): + def __create_effnet_encoder_components(self, frame, row: int, ui_state) -> int: # effnet encoder model - components.label(frame, row, 0, "Effnet Encoder Model", + self.components.label(frame, row, 0, "Effnet Encoder Model", tooltip="Filename, directory or Hugging Face repository of the effnet encoder model") - components.path_entry( - frame, row, 1, self.ui_state, "effnet_encoder.model_name", - mode="file", path_modifier=components.json_path_modifier + self.components.path_entry( + frame, row, 1, ui_state, "effnet_encoder.model_name", + mode="file", path_modifier=path_util.json_path_modifier ) # effnet encoder weight dtype - components.label(frame, row, 3, "Effnet Encoder Data Type", + self.components.label(frame, row, 3, "Effnet Encoder Data Type", tooltip="The effnet encoder weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "effnet_encoder.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "effnet_encoder.weight_dtype") row += 1 @@ -588,38 +600,39 @@ def __create_decoder_components( self, frame, row: int, + ui_state, has_text_encoder: bool, ) -> int: # decoder model - components.label(frame, row, 0, "Decoder Model", + self.components.label(frame, row, 0, "Decoder Model", tooltip="Filename, directory or Hugging Face repository of the decoder model") - components.path_entry( - frame, row, 1, self.ui_state, "decoder.model_name", - mode="file", path_modifier=components.json_path_modifier + self.components.path_entry( + frame, row, 1, ui_state, "decoder.model_name", + mode="file", path_modifier=path_util.json_path_modifier ) # decoder weight dtype - components.label(frame, row, 3, "Decoder Data Type", + self.components.label(frame, row, 3, "Decoder Data Type", tooltip="The decoder weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "decoder.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "decoder.weight_dtype") row += 1 if has_text_encoder: # decoder text encoder weight dtype - components.label(frame, row, 3, "Decoder Text Encoder Data Type", + self.components.label(frame, row, 3, "Decoder Text Encoder Data Type", tooltip="The decoder text encoder weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "decoder_text_encoder.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "decoder_text_encoder.weight_dtype") row += 1 # decoder vqgan weight dtype - components.label(frame, row, 3, "Decoder VQGAN Data Type", + self.components.label(frame, row, 3, "Decoder VQGAN Data Type", tooltip="The decoder vqgan weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "decoder_vqgan.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "decoder_vqgan.weight_dtype") row += 1 @@ -629,30 +642,31 @@ def __create_output_components( self, frame, row: int, + ui_state, allow_safetensors: bool = False, allow_diffusers: bool = False, allow_legacy_safetensors: bool = False, allow_comfy: bool = False, ) -> int: # output model destination - components.label(frame, row, 0, "Model Output Destination", + self.components.label(frame, row, 0, "Model Output Destination", tooltip="Filename or directory where the output model is saved") - components.path_entry( - frame, row, 1, self.ui_state, "output_model_destination", + self.components.path_entry( + frame, row, 1, ui_state, "output_model_destination", mode="file", io_type=PathIOType.MODEL, ) # output data type - components.label(frame, row, 3, "Output Data Type", + self.components.label(frame, row, 3, "Output Data Type", tooltip="Precision to use when saving the output model") - components.options_kv(frame, row, 4, [ + self.components.options_kv(frame, row, 4, [ ("float16", DataType.FLOAT_16), ("float32", DataType.FLOAT_32), ("bfloat16", DataType.BFLOAT_16), ("float8", DataType.FLOAT_8), ("nfloat4", DataType.NFLOAT_4), - ], self.ui_state, "output_dtype") + ], ui_state, "output_dtype") row += 1 @@ -667,21 +681,21 @@ def __create_output_components( if allow_comfy: formats.append(("Comfy LoRA", ModelFormat.COMFY_LORA)) - components.label(frame, row, 0, "Output Format", + self.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") + self.components.options_kv(frame, row, 1, formats, ui_state, "output_model_format") # include config - components.label(frame, row, 3, "Include Config", + self.components.label(frame, row, 3, "Include Config", tooltip="Include the training configuration in the final model. Only supported for safetensors files. " "None: No config is included. " "Settings: All training settings are included. " "All: All settings, including the samples and concepts are included.") - components.options_kv(frame, row, 4, [ + self.components.options_kv(frame, row, 4, [ ("None", ConfigPart.NONE), ("Settings", ConfigPart.SETTINGS), ("All", ConfigPart.ALL), - ], self.ui_state, "include_train_config") + ], ui_state, "include_train_config") row += 1 diff --git a/modules/ui/BaseMuonAdamWindowView.py b/modules/ui/BaseMuonAdamWindowView.py index 5879ab432..90dc04eee 100644 --- a/modules/ui/BaseMuonAdamWindowView.py +++ b/modules/ui/BaseMuonAdamWindowView.py @@ -1,67 +1,10 @@ -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.Optimizer import Optimizer -from modules.util.optimizer_util import OPTIMIZER_DEFAULT_PARAMETERS -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState -import customtkinter as ctk -MUON_AUX_ADAM_DEFAULTS = { - "beta1": 0.9, - "beta2": 0.999, - "eps": 1e-8, - "weight_decay": 0.0, -} +class BaseMuonAdamWindowView: + def __init__(self, components): + self.components = components -class MuonAdamWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - train_config: TrainConfig, - ui_state: UIState, - parent_optimizer_type: Optimizer, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) - - self.parent = parent - self.train_config = train_config - self.adam_ui_state = ui_state - self.parent_optimizer_type = parent_optimizer_type - - if self.parent_optimizer_type == Optimizer.MUON: - self.title("Muon's Auxiliary AdamW Settings") - self.adam_params_def = MUON_AUX_ADAM_DEFAULTS - else: - self.title("Muon_adv's Auxiliary AdamW_adv Settings") - self.adam_params_def = OPTIMIZER_DEFAULT_PARAMETERS[Optimizer.ADAMW_ADV] - - self.geometry("800x500") - self.resizable(True, True) - - self.grid_rowconfigure(0, weight=1) - self.grid_rowconfigure(1, weight=0) - self.grid_columnconfigure(0, weight=1) - - self.frame = ctk.CTkScrollableFrame(self, fg_color="transparent") - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) - self.frame.grid_columnconfigure(2, minsize=50) - self.frame.grid_columnconfigure(3, weight=0) - self.frame.grid_columnconfigure(4, weight=1) - - components.button(self, 1, 0, "ok", command=self.destroy) - self.create_adam_params_ui(self.frame) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - def create_adam_params_ui(self, master): + def build_content(self, master, controller, ui_state): # This is a large map, copied from OptimizerParamsWindow for simplicity. # @formatter:off KEY_DETAIL_MAP = { @@ -84,7 +27,7 @@ def create_adam_params_ui(self, master): } # @formatter:on - adam_params = self.adam_params_def + adam_params = controller.get_adam_params_def() for index, key in enumerate(adam_params.keys()): if key not in KEY_DETAIL_MAP: @@ -99,9 +42,9 @@ def create_adam_params_ui(self, master): row = index // 2 col = 3 * (index % 2) - components.label(master, row, col, title, tooltip=tooltip) + self.components.label(master, row, col, title, tooltip=tooltip) if param_type != 'bool': - components.entry(master, row, col + 1, self.adam_ui_state, key) + self.components.entry(master, row, col + 1, ui_state, key) else: - components.switch(master, row, col + 1, self.adam_ui_state, key) + self.components.switch(master, row, col + 1, ui_state, key) diff --git a/modules/ui/BaseOffloadingWindowView.py b/modules/ui/BaseOffloadingWindowView.py index 54035e121..8ec9c0cb1 100644 --- a/modules/ui/BaseOffloadingWindowView.py +++ b/modules/ui/BaseOffloadingWindowView.py @@ -1,75 +1,24 @@ -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.GradientCheckpointingMethod import ( - GradientCheckpointingMethod, -) -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState +from modules.util.enum.GradientCheckpointingMethod import GradientCheckpointingMethod -import customtkinter as ctk +class BaseOffloadingWindowView: + def __init__(self, components): + self.components = components -class OffloadingWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - config: TrainConfig, - ui_state: UIState, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) + def build_content(self, frame, controller, ui_state): + self.components.label(frame, 0, 0, "Gradient checkpointing", + tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") + self.components.options(frame, 0, 1, [str(x) for x in list(GradientCheckpointingMethod)], ui_state, + "gradient_checkpointing") - self.config = config - self.ui_state = ui_state - self.image_preview_file_index = 0 - self.ax = None - self.canvas = None + self.components.label(frame, 1, 0, "Async Offloading", + tooltip="Enables Asynchronous offloading.") + self.components.switch(frame, 1, 1, ui_state, "enable_async_offloading") - self.title("Offloading") - self.geometry("800x400") - self.resizable(True, True) + self.components.label(frame, 2, 0, "Offload Activations", + tooltip="Enables Activation Offloading") + self.components.switch(frame, 2, 1, ui_state, "enable_activation_offloading") - self.grid_rowconfigure(0, weight=1) - self.grid_columnconfigure(0, weight=1) - - frame = self.__content_frame(self) - frame.grid(row=0, column=0, sticky='nsew') - components.button(self, 1, 0, "ok", self.__ok) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def __content_frame(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=1) - frame.grid_columnconfigure(1, weight=1) - - # timestep distribution - components.label(frame, 0, 0, "Gradient checkpointing", - tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") - components.options(frame, 0, 1, [str(x) for x in list(GradientCheckpointingMethod)], self.ui_state, - "gradient_checkpointing") - - # gradient checkpointing layer offloading - components.label(frame, 1, 0, "Async Offloading", - tooltip="Enables Asynchronous offloading.") - components.switch(frame, 1, 1, self.ui_state, "enable_async_offloading") - - # gradient checkpointing layer offloading - components.label(frame, 2, 0, "Offload Activations", - tooltip="Enables Activation Offloading") - components.switch(frame, 2, 1, self.ui_state, "enable_activation_offloading") - - # gradient checkpointing layer offloading - components.label(frame, 3, 0, "Layer offload fraction", - tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") - components.entry(frame, 3, 1, self.ui_state, "layer_offload_fraction") - - frame.pack(fill="both", expand=1) - return frame - - def __ok(self): - self.destroy() + self.components.label(frame, 3, 0, "Layer offload fraction", + tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") + self.components.entry(frame, 3, 1, ui_state, "layer_offload_fraction") diff --git a/modules/ui/BaseOptimizerParamsWindowView.py b/modules/ui/BaseOptimizerParamsWindowView.py index 16063c26c..b5a5c8911 100644 --- a/modules/ui/BaseOptimizerParamsWindowView.py +++ b/modules/ui/BaseOptimizerParamsWindowView.py @@ -1,94 +1,32 @@ -import contextlib -from tkinter import TclError -from modules.ui.MuonAdamWindow import MUON_AUX_ADAM_DEFAULTS, MuonAdamWindow -from modules.util.config.TrainConfig import TrainConfig, TrainOptimizerConfig from modules.util.enum.Optimizer import Optimizer from modules.util.optimizer_util import ( OPTIMIZER_DEFAULT_PARAMETERS, - change_optimizer, - load_optimizer_defaults, - update_optimizer_config, ) -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState -import customtkinter as ctk +class BaseOptimizerParamsWindowView: + def __init__(self, components): + self.components = components -class OptimizerParamsWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - train_config: TrainConfig, - ui_state, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) - - self.parent = parent - self.train_config = train_config - self.ui_state = ui_state - self.optimizer_ui_state = ui_state.get_var("optimizer") - self.protocol("WM_DELETE_WINDOW", self.on_window_close) - self.muon_adam_button = None - - self.title("Optimizer Settings") - self.geometry("800x500") - self.resizable(True, True) - - self.grid_rowconfigure(0, weight=1) - self.grid_rowconfigure(1, weight=0) - self.grid_columnconfigure(0, weight=1) - - self.frame = ctk.CTkScrollableFrame(self, fg_color="transparent") - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) - self.frame.grid_columnconfigure(2, minsize=50) - self.frame.grid_columnconfigure(3, weight=0) - self.frame.grid_columnconfigure(4, weight=1) - - components.button(self, 1, 0, "ok", command=self.on_window_close) - self.main_frame(self.frame) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def main_frame(self, master): + def build_content(self, frame, controller, ui_state, optimizer_ui_state, + on_optimizer_change_cb, load_defaults_cb): # Optimizer - components.label(master, 0, 0, "Optimizer", - tooltip="The type of optimizer") + self.components.label(frame, 0, 0, "Optimizer", + tooltip="The type of optimizer") # Create the optimizer dropdown menu and set the command - components.options(master, 0, 1, [str(x) for x in list(Optimizer)], self.optimizer_ui_state, "optimizer", - command=self.on_optimizer_change) + self.components.options(frame, 0, 1, [str(x) for x in list(Optimizer)], optimizer_ui_state, "optimizer", + command=on_optimizer_change_cb) # Defaults Button - components.label(master, 0, 3, "Optimizer Defaults", - tooltip="Load default settings for the selected optimizer") - components.button(self.frame, 0, 4, "Load Defaults", self.load_defaults, - tooltip="Load default settings for the selected optimizer") - - self.create_dynamic_ui(master) - - def clear_dynamic_ui(self, master): - with contextlib.suppress(TclError): - for widget in master.winfo_children(): - grid_info = widget.grid_info() - if int(grid_info["row"]) >= 1: - widget.destroy() - - def create_dynamic_ui( - self, - master, - ): + self.components.label(frame, 0, 3, "Optimizer Defaults", + tooltip="Load default settings for the selected optimizer") + self.components.button(frame, 0, 4, "Load Defaults", load_defaults_cb, + tooltip="Load default settings for the selected optimizer") + def build_dynamic_content(self, master, controller, optimizer_ui_state, + update_user_pref_cb, open_muon_adam_cb): # Lookup for the title and tooltip for a key # @formatter:off KEY_DETAIL_MAP = { @@ -197,10 +135,7 @@ def create_dynamic_ui( } # @formatter:on - if not self.winfo_exists(): # check if this window isn't open - return - - selected_optimizer = self.train_config.optimizer.optimizer + selected_optimizer = controller.config.optimizer.optimizer # Extract the keys for the selected optimizer for index, key in enumerate(OPTIMIZER_DEFAULT_PARAMETERS[selected_optimizer].keys()): @@ -215,74 +150,20 @@ def create_dynamic_ui( row = (index // 2) + 1 col = 3 * (index % 2) - components.label(master, row, col, title, tooltip=tooltip) + self.components.label(master, row, col, title, tooltip=tooltip) if key == 'MuonWithAuxAdam': - frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=row, column=col + 1, columnspan=2, sticky="ew", padx=0, pady=0) - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) + frame = self.components.inline_frame(master, row, col + 1, columnspan=2) - components.switch(frame, 0, 0, self.optimizer_ui_state, key, command=self.update_user_pref) + self.components.switch(frame, 0, 0, optimizer_ui_state, key, command=update_user_pref_cb) - self.muon_adam_button = components.button( - frame, 0, 1, "...", self.open_muon_adam_window, + self.muon_adam_button = self.components.button( + frame, 0, 1, "...", open_muon_adam_cb, tooltip="Configure the auxiliary AdamW_adv optimizer", - width=20, padx=5 ) - self.toggle_muon_adam_button() + width=20, padx=5) elif type != 'bool': - components.entry(master, row, col + 1, self.optimizer_ui_state, key, - command=self.update_user_pref) + self.components.entry(master, row, col + 1, optimizer_ui_state, key, + command=update_user_pref_cb) else: - components.switch(master, row, col + 1, self.optimizer_ui_state, key, - command=self.update_user_pref) - - def update_user_pref(self, *args): - update_optimizer_config(self.train_config) - self.toggle_muon_adam_button() - - def on_optimizer_change(self, *args): - optimizer_config = change_optimizer(self.train_config) - self.ui_state.get_var("optimizer").update(optimizer_config) - - self.clear_dynamic_ui(self.frame) - self.create_dynamic_ui(self.frame) - - def load_defaults(self, *args): - optimizer_config = load_optimizer_defaults(self.train_config) - self.ui_state.get_var("optimizer").update(optimizer_config) - - def on_window_close(self): - self.destroy() - - def toggle_muon_adam_button(self): - if self.muon_adam_button and self.muon_adam_button.winfo_exists(): - muon_with_adam = self.optimizer_ui_state.get_var("MuonWithAuxAdam").get() - self.muon_adam_button.configure(state="normal" if muon_with_adam else "disabled") - - def open_muon_adam_window(self): - current_optimizer = self.train_config.optimizer.optimizer - - adam_config = TrainOptimizerConfig.default_values() - current_state = self.train_config.optimizer.muon_adam_config - - if current_optimizer == Optimizer.MUON: - defaults = MUON_AUX_ADAM_DEFAULTS - else: - defaults = OPTIMIZER_DEFAULT_PARAMETERS[Optimizer.ADAMW_ADV] - - if current_state is None: - adam_config.from_dict(defaults) - if current_optimizer != Optimizer.MUON: - adam_config.optimizer = Optimizer.ADAMW_ADV - elif isinstance(current_state, dict): - adam_config.from_dict(current_state) - else: - # Should not happen if TrainConfig defines it as dict, but for safety - adam_config = current_state - - temp_adam_ui_state = UIState(self, adam_config) - window = MuonAdamWindow(self, self.train_config, temp_adam_ui_state, current_optimizer) - self.wait_window(window) - - self.train_config.optimizer.muon_adam_config = adam_config.to_dict() + self.components.switch(master, row, col + 1, optimizer_ui_state, key, + command=update_user_pref_cb) diff --git a/modules/ui/BaseProfilingWindowView.py b/modules/ui/BaseProfilingWindowView.py index 8d298abe3..8e0de3b64 100644 --- a/modules/ui/BaseProfilingWindowView.py +++ b/modules/ui/BaseProfilingWindowView.py @@ -1,57 +1,21 @@ -import faulthandler +from abc import abstractmethod -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -import customtkinter as ctk -from scalene import scalene_profiler +class BaseProfilingWindowView: + def __init__(self, components): + self.components = components - -class ProfilingWindow(ctk.CTkToplevel): - def __init__(self, parent, *args, **kwargs): - super().__init__(parent, *args, **kwargs) - self.parent = parent - - self.title("Profiling") - self.geometry("512x512") - self.resizable(True, True) - self.wait_visibility() - self.focus_set() - - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=0) - self.grid_rowconfigure(2, weight=1) - self.grid_columnconfigure(0, weight=1) - - components.button(self, 0, 0, "Dump stack", self._dump_stack) - self._profile_button = components.button( - self, 1, 0, "Start Profiling", self._start_profiler, + def build_content(self, frame, bottom_bar, controller): + self.components.button(frame, 0, 0, "Dump stack", controller.dump_stack) + self._profile_button = self.components.button( + frame, 1, 0, "Start Profiling", controller.start_profiler, tooltip="Turns on/off Scalene profiling. Only works when OneTrainer is launched with Scalene!") + self._message_label = self.components.label(bottom_bar, 0, 0, "Inactive") - # Bottom bar - self._bottom_bar = ctk.CTkFrame(master=self, corner_radius=0) - self._bottom_bar.grid(row=2, column=0, sticky="sew") - self._message_label = components.label(self._bottom_bar, 0, 0, "Inactive") - - self.protocol("WM_DELETE_WINDOW", self.withdraw) - self.withdraw() - self.after(200, lambda: set_window_icon(self)) - - def _dump_stack(self): - with open('stacks.txt', 'w') as f: - faulthandler.dump_traceback(f) - self._message_label.configure(text='Stack dumped to stacks.txt') - - def _end_profiler(self): - scalene_profiler.stop() - - self._message_label.configure(text='Inactive') - self._profile_button.configure(text='Start Profiling') - self._profile_button.configure(command=self._start_profiler) - - def _start_profiler(self): - scalene_profiler.start() + @abstractmethod + def set_message(self, text): + pass - self._message_label.configure(text='Profiling active...') - self._profile_button.configure(text='End Profiling') - self._profile_button.configure(command=self._end_profiler) + @abstractmethod + def set_profiling_active(self, active): + pass diff --git a/modules/ui/BaseSampleFrameView.py b/modules/ui/BaseSampleFrameView.py index 297caac29..c3e5860c9 100644 --- a/modules/ui/BaseSampleFrameView.py +++ b/modules/ui/BaseSampleFrameView.py @@ -1,98 +1,59 @@ -from modules.util.config.SampleConfig import SampleConfig -from modules.util.enum.ModelType import ModelType from modules.util.enum.NoiseScheduler import NoiseScheduler -from modules.util.ui import components -from modules.util.ui.UIState import UIState -import customtkinter as ctk +class BaseSampleFrameView: + def __init__(self, components): + self.components = components -class SampleFrame(ctk.CTkFrame): - def __init__( - self, - parent, - sample: SampleConfig, - ui_state: UIState, - model_type: ModelType, - include_prompt: bool = True, - include_settings: bool = True, - ): - ctk.CTkFrame.__init__(self, parent, fg_color="transparent") - - self.sample = sample - self.ui_state = ui_state - self.model_type = model_type - - is_flow_matching = model_type.is_flow_matching() - is_inpainting_model = model_type.has_conditioning_image_input() - is_video_model = model_type.is_video_model() - - if include_prompt and include_prompt: - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=1) - self.grid_columnconfigure(0, weight=1) - - if include_prompt: - top_frame = ctk.CTkFrame(self, fg_color="transparent") - top_frame.grid(row=0, column=0, padx=0, pady=0, sticky="nsew") - - top_frame.grid_columnconfigure(0, weight=0) - top_frame.grid_columnconfigure(1, weight=1) - - if include_settings: - bottom_frame = ctk.CTkFrame(self, fg_color="transparent") - bottom_frame.grid(row=1, column=0, padx=0, pady=0, sticky="nsew") - - bottom_frame.grid_columnconfigure(0, weight=0) - bottom_frame.grid_columnconfigure(1, weight=1) - bottom_frame.grid_columnconfigure(2, weight=0) - bottom_frame.grid_columnconfigure(3, weight=1) - + def build_content(self, top_frame, bottom_frame, ui_state, controller, include_prompt, include_settings): + is_flow_matching = controller.is_flow_matching() + is_inpainting_model = controller.is_inpainting_model() + is_video_model = controller.is_video_model() if include_prompt: # prompt - components.label(top_frame, 0, 0, "prompt:") - components.entry(top_frame, 0, 1, self.ui_state, "prompt") + self.components.label(top_frame, 0, 0, "prompt:") + self.components.entry(top_frame, 0, 1, ui_state, "prompt") # negative prompt - components.label(top_frame, 1, 0, "negative prompt:") - components.entry(top_frame, 1, 1, self.ui_state, "negative_prompt") + self.components.label(top_frame, 1, 0, "negative prompt:") + self.components.entry(top_frame, 1, 1, ui_state, "negative_prompt") if include_settings: # width - components.label(bottom_frame, 0, 0, "width:") - components.entry(bottom_frame, 0, 1, self.ui_state, "width") + self.components.label(bottom_frame, 0, 0, "width:") + self.components.entry(bottom_frame, 0, 1, ui_state, "width") # height - components.label(bottom_frame, 0, 2, "height:") - components.entry(bottom_frame, 0, 3, self.ui_state, "height") + self.components.label(bottom_frame, 0, 2, "height:") + self.components.entry(bottom_frame, 0, 3, ui_state, "height") if is_video_model: # frames - components.label(bottom_frame, 1, 0, "frames:", - tooltip="Number of frames to generate. Only used when generating videos.") - components.entry(bottom_frame, 1, 1, self.ui_state, "frames") + self.components.label(bottom_frame, 1, 0, "frames:", + tooltip="Number of frames to generate. Only used when generating videos.") + self.components.entry(bottom_frame, 1, 1, ui_state, "frames") # length - components.label(bottom_frame, 1, 2, "length:", - tooltip="Length in seconds of audio output.") - components.entry(bottom_frame, 1, 3, self.ui_state, "length") + self.components.label(bottom_frame, 1, 2, "length:", + tooltip="Length in seconds of audio output.") + self.components.entry(bottom_frame, 1, 3, ui_state, "length") # seed - components.label(bottom_frame, 2, 0, "seed:") - components.entry(bottom_frame, 2, 1, self.ui_state, "seed") + self.components.label(bottom_frame, 2, 0, "seed:") + self.components.entry(bottom_frame, 2, 1, ui_state, "seed") # random seed - components.label(bottom_frame, 2, 2, "random seed:") - components.switch(bottom_frame, 2, 3, self.ui_state, "random_seed") + self.components.label(bottom_frame, 2, 2, "random seed:") + self.components.switch(bottom_frame, 2, 3, ui_state, "random_seed") # cfg scale - components.label(bottom_frame, 3, 0, "cfg scale:") - components.entry(bottom_frame, 3, 1, self.ui_state, "cfg_scale") + self.components.label(bottom_frame, 3, 0, "cfg scale:") + self.components.entry(bottom_frame, 3, 1, ui_state, "cfg_scale") # sampler if not is_flow_matching: - components.label(bottom_frame, 4, 2, "sampler:") - components.options_kv(bottom_frame, 4, 3, [ + self.components.label(bottom_frame, 4, 2, "sampler:") + self.components.options_kv(bottom_frame, 4, 3, [ ("DDIM", NoiseScheduler.DDIM), ("Euler", NoiseScheduler.EULER), ("Euler A", NoiseScheduler.EULER_A), @@ -103,32 +64,32 @@ def __init__( ("DPM++ Karras", NoiseScheduler.DPMPP_KARRAS), ("DPM++ SDE Karras", NoiseScheduler.DPMPP_SDE_KARRAS), ("UniPC Karras", NoiseScheduler.UNIPC_KARRAS) - ], self.ui_state, "noise_scheduler") + ], ui_state, "noise_scheduler") # steps - components.label(bottom_frame, 4, 0, "steps:") - components.entry(bottom_frame, 4, 1, self.ui_state, "diffusion_steps") + self.components.label(bottom_frame, 4, 0, "steps:") + self.components.entry(bottom_frame, 4, 1, ui_state, "diffusion_steps") # inpainting if is_inpainting_model: - components.label(bottom_frame, 5, 0, "inpainting:", - tooltip="Enables inpainting sampling. Only available when sampling from an inpainting model.") - components.switch(bottom_frame, 5, 1, self.ui_state, "sample_inpainting") + self.components.label(bottom_frame, 5, 0, "inpainting:", + tooltip="Enables inpainting sampling. Only available when sampling from an inpainting model.") + self.components.switch(bottom_frame, 5, 1, ui_state, "sample_inpainting") # 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", - mode="file", - allow_model_files=False, - allow_image_files=True, - ) + self.components.label(bottom_frame, 6, 0, "base image path:", + tooltip="The base image used when inpainting.") + self.components.path_entry(bottom_frame, 6, 1, ui_state, "base_image_path", + mode="file", + allow_model_files=False, + allow_image_files=True, + ) # 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", - mode="file", - allow_model_files=False, - allow_image_files=True, - ) + self.components.label(bottom_frame, 6, 2, "mask image path:", + tooltip="The mask used when inpainting.") + self.components.path_entry(bottom_frame, 6, 3, ui_state, "mask_image_path", + mode="file", + allow_model_files=False, + allow_image_files=True, + ) diff --git a/modules/ui/BaseSampleParamsWindowView.py b/modules/ui/BaseSampleParamsWindowView.py index 2b0b3f3f1..bbe410a9a 100644 --- a/modules/ui/BaseSampleParamsWindowView.py +++ b/modules/ui/BaseSampleParamsWindowView.py @@ -1,39 +1,6 @@ -from modules.ui.SampleFrame import SampleFrame -from modules.util.config.SampleConfig import SampleConfig -from modules.util.enum.ModelType import ModelType -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState -import customtkinter as ctk -class SampleParamsWindow(ctk.CTkToplevel): - def __init__(self, parent, sample: SampleConfig, ui_state: UIState, model_type: ModelType | None = None, *args, **kwargs): - super().__init__(parent, *args, **kwargs) - - self.sample = sample - self.ui_state = ui_state - self.model_type = model_type - - self.title("Sample") - self.geometry("800x500") - self.resizable(True, True) - - self.grid_rowconfigure(0, weight=1) - self.grid_rowconfigure(1, weight=0) - self.grid_columnconfigure(0, weight=1) - - frame = SampleFrame(self, self.sample, self.ui_state, model_type=model_type) - frame.grid(row=0, column=0, padx=0, pady=0, sticky="nsew") - - components.button(self, 1, 0, "ok", self.__ok) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def __ok(self): - self.destroy() +class BaseSampleParamsWindowView: + def __init__(self, components): + self.components = components diff --git a/modules/ui/BaseSampleWindowView.py b/modules/ui/BaseSampleWindowView.py index 0f91ad2fa..1a9c6c0c6 100644 --- a/modules/ui/BaseSampleWindowView.py +++ b/modules/ui/BaseSampleWindowView.py @@ -1,227 +1,8 @@ -import contextlib -import copy -import os -import tkinter as tk -import traceback -from modules.model.BaseModel import BaseModel -from modules.modelSampler.BaseModelSampler import ( - BaseModelSampler, - ModelSamplerOutput, -) -from modules.ui.SampleFrame import SampleFrame -from modules.util import create -from modules.util.callbacks.TrainCallbacks import TrainCallbacks -from modules.util.commands.TrainCommands import TrainCommands -from modules.util.config.SampleConfig import SampleConfig -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.EMAMode import EMAMode -from modules.util.enum.FileType import FileType -from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.time_util import get_string_timestamp -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState -import torch -import customtkinter as ctk -from PIL import Image -class SampleWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - train_config: TrainConfig, - use_external_model: bool, - callbacks: TrainCallbacks | None = None, - commands: TrainCommands | None = None, - *args, **kwargs - ): - super().__init__(parent, *args, **kwargs) - - self.title("Sample") - self.geometry("1200x800") - self.resizable(True, True) - - if not use_external_model: - self.initial_train_config = TrainConfig.default_values().from_dict(train_config.to_dict()) - # remove some settings to speed up model loading for sampling - self.initial_train_config.optimizer.optimizer = None - self.initial_train_config.ema = EMAMode.OFF - else: - self.initial_train_config = None - - #TODO why is there a current_train_config and an initial_train_config? - #current_train_config doesn't seem to ever change - self.current_train_config = train_config - self.callbacks = callbacks - self.commands = commands - - # get model specific defaults - model_type = train_config.model_type - self.sample = SampleConfig.default_values(model_type) - self.ui_state = UIState(self, self.sample) - - if use_external_model: - self.callbacks.set_on_sample_custom(self.__update_preview) - self.callbacks.set_on_update_sample_custom_progress(self.__update_progress) - else: - self.model = None - self.model_sampler = None - - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=1) - self.grid_rowconfigure(2, weight=0) - self.grid_rowconfigure(3, weight=0) - self.grid_columnconfigure(0, weight=0) - self.grid_columnconfigure(1, weight=1) - - prompt_frame = SampleFrame(self, self.sample, self.ui_state, include_settings=False, model_type=model_type) - prompt_frame.grid(row=0, column=0, columnspan=2, padx=0, pady=0, sticky="nsew") - - settings_frame = SampleFrame(self, self.sample, self.ui_state, include_prompt=False, model_type=model_type) - settings_frame.grid(row=1, column=0, padx=0, pady=0, sticky="nsew") - - # image - self.image = ctk.CTkImage( - light_image=self.__dummy_image(), - size=(512, 512) - ) - - image_label = ctk.CTkLabel(master=self, text="", image=self.image, height=512, width=512) - image_label.grid(row=1, column=1, rowspan=3, sticky="nsew") - - self.progress = components.progress(self, 2, 0) - components.button(self, 3, 0, "sample", self.__sample) - - self.wait_visibility() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - def __load_model(self) -> BaseModel: - model_loader = create.create_model_loader( - model_type=self.initial_train_config.model_type, - training_method=self.initial_train_config.training_method, - ) - - model_setup = create.create_model_setup( - model_type=self.initial_train_config.model_type, - train_device=torch.device(self.initial_train_config.train_device), - temp_device=torch.device(self.initial_train_config.temp_device), - training_method=self.initial_train_config.training_method, - ) - - model_names = self.initial_train_config.model_names() - if self.initial_train_config.continue_last_backup: - last_backup_path = self.initial_train_config.get_last_backup_path() - - if last_backup_path: - if self.initial_train_config.training_method == TrainingMethod.LORA: - model_names.lora = last_backup_path - elif self.initial_train_config.training_method == TrainingMethod.EMBEDDING: - model_names.embedding.model_name = last_backup_path - else: # fine-tunes - model_names.base_model = last_backup_path - - print(f"Loading from backup '{last_backup_path}'...") - else: - print("No backup found, loading without backup...") - - if self.initial_train_config.quantization.cache_dir is None: - self.initial_train_config.quantization.cache_dir = self.initial_train_config.cache_dir + "/quantization" - os.makedirs(self.initial_train_config.quantization.cache_dir, exist_ok=True) - - model = model_loader.load( - model_type=self.initial_train_config.model_type, - model_names=model_names, - weight_dtypes=self.initial_train_config.weight_dtypes(), - quantization=self.initial_train_config.quantization, - ) - model.train_config = self.initial_train_config - - model_setup.setup_optimizations(model, self.initial_train_config) - model_setup.setup_train_device(model, self.initial_train_config) - model_setup.setup_model(model, self.initial_train_config) - model.to(torch.device(self.initial_train_config.temp_device)) - - return model - - def __create_sampler(self, model: BaseModel) -> BaseModelSampler: - return create.create_model_sampler( - train_device=torch.device(self.initial_train_config.train_device), - temp_device=torch.device(self.initial_train_config.temp_device), - model=model, - model_type=self.initial_train_config.model_type, - training_method=self.initial_train_config.training_method, - ) - - def __update_preview(self, sampler_output: ModelSamplerOutput): - if sampler_output.file_type == FileType.IMAGE: - image = sampler_output.data - self.image.configure( - light_image=image, - size=(image.width, image.height), - ) - - def __update_progress(self, progress: int, max_progress: int): - self.progress.set(progress / max_progress) - self.update() - - def __dummy_image(self) -> Image: - return Image.new(mode="RGB", size=(512, 512), color=(0, 0, 0)) - - def __sample(self): - sample = copy.copy(self.sample) - - if self.commands: - self.commands.sample_custom(sample) - else: - if self.model is None: - # lazy initialization - self.model = self.__load_model() - self.model_sampler = self.__create_sampler(self.model) - - sample.from_train_config(self.current_train_config) - - sample_dir = os.path.join( - self.initial_train_config.workspace_dir, - "samples", - "custom", - ) - - progress = self.model.train_progress - sample_path = os.path.join( - sample_dir, - f"{get_string_timestamp()}-training-sample-{progress.filename_string()}" - ) - - self.model.eval() - - self.model_sampler.sample( - sample_config=sample, - destination=sample_path, - image_format=self.current_train_config.sample_image_format, - video_format=self.current_train_config.sample_video_format, - audio_format=self.current_train_config.sample_audio_format, - on_sample=self.__update_preview, - on_update_progress=self.__update_progress, - ) - - def destroy(self): - try: - if hasattr(self, "_icon_image_ref"): - del self._icon_image_ref - - # Remove any pending after callbacks - for after_id in self.tk.call('after', 'info'): - with contextlib.suppress(tk.TclError, RuntimeError): - self.after_cancel(after_id) - - super().destroy() - except (tk.TclError, RuntimeError) as e: - print(f"Error destroying window: {e}") - except Exception as e: - print(f"Unexpected error destroying window: {e}") - traceback.print_exc() +class BaseSampleWindowView: + def __init__(self, components): + pass diff --git a/modules/ui/BaseSamplingTabView.py b/modules/ui/BaseSamplingTabView.py index 5a3c44f08..69ec082a1 100644 --- a/modules/ui/BaseSamplingTabView.py +++ b/modules/ui/BaseSamplingTabView.py @@ -1,124 +1,67 @@ -from modules.ui.ConfigList import ConfigList -from modules.ui.SampleParamsWindow import SampleParamsWindow -from modules.util.config.SampleConfig import SampleConfig -from modules.util.config.TrainConfig import TrainConfig -from modules.util.ui import components -from modules.util.ui.UIState import UIState +from abc import ABC, abstractmethod -import customtkinter as ctk +from modules.ui.BaseConfigListView import BaseConfigListView -class SamplingTab(ConfigList): +class BaseSamplingTabView(BaseConfigListView): + pass - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__( - master, - train_config, - ui_state, - from_external_file=True, - attr_name="sample_definition_file_name", - config_dir="training_samples", - default_config_name="samples.json", - add_button_text="Add Sample", - add_button_tooltip="Add a new sample configuration.", - is_full_width=True, - show_toggle_button=True - ) - def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - return SampleWidget(master, element, i, open_command, remove_command, clone_command, save_command) - - def create_new_element(self) -> dict: - return SampleConfig.default_values(self.train_config.model_type) - - def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: - return SampleParamsWindow(self.master, self.current_config[i], ui_state, model_type=self.train_config.model_type) - - -class SampleWidget(ctk.CTkFrame): - def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): - super().__init__( - master=master, corner_radius=10, bg_color="transparent" - ) +class BaseSampleWidgetView(ABC): + def __init__(self, components): + self.components = components + def build_content(self, frame, element, ui_state, i, open_command, remove_command, clone_command, save_command): self.element = element - self.ui_state = UIState(self, element) self.i = i self.save_command = save_command - self.grid_columnconfigure(10, weight=1) - # close button - close_button = ctk.CTkButton( - master=self, - width=20, - height=20, - text="X", - corner_radius=2, - fg_color="#C00000", - command=lambda: remove_command(self.i), - ) - close_button.grid(row=0, column=0) + self.components.colored_icon_button(frame, 0, 0, "X", "#C00000", lambda: remove_command(self.i)) # clone button - clone_button = ctk.CTkButton( - master=self, - width=20, - height=20, - text="+", - corner_radius=2, - fg_color="#00C000", - command=lambda: clone_command(self.i), - ) - clone_button.grid(row=0, column=1, padx=5) + self.components.colored_icon_button(frame, 0, 1, "+", "#00C000", lambda: clone_command(self.i), padx=5) # enabled - self.enabled_switch = components.switch(self, 0, 2, self.ui_state, "enabled", self.__switch_enabled) - self.enabled_switch.configure(width=40) + self.enabled_switch = self.components.switch(frame, 0, 2, ui_state, "enabled", self._switch_enabled, width=40) # width - components.label(self, 0, 3, "width:") - self.width_entry = components.entry(self, 0, 4, self.ui_state, "width") - self.width_entry.bind('', lambda _: save_command()) - self.width_entry.configure(width=50) + self.components.label(frame, 0, 3, "width:") + self.width_entry = self.components.entry(frame, 0, 4, ui_state, "width", width=50) # height - components.label(self, 0, 5, "height:") - self.height_entry = components.entry(self, 0, 6, self.ui_state, "height") - self.height_entry.bind('', lambda _: save_command()) - self.height_entry.configure(width=50) + self.components.label(frame, 0, 5, "height:") + self.height_entry = self.components.entry(frame, 0, 6, ui_state, "height", width=50) # seed - components.label(self, 0, 7, "seed:") - self.seed_entry = components.entry(self, 0, 8, self.ui_state, "seed") - self.seed_entry.bind('', lambda _: save_command()) - self.seed_entry.configure(width=80) + self.components.label(frame, 0, 7, "seed:") + self.seed_entry = self.components.entry(frame, 0, 8, ui_state, "seed", width=80) # prompt - components.label(self, 0, 9, "prompt:") - self.prompt_entry = components.entry(self, 0, 10, self.ui_state, "prompt") - self.prompt_entry.bind('', lambda _: save_command()) + self.components.label(frame, 0, 9, "prompt:") + self.prompt_entry = self.components.entry(frame, 0, 10, ui_state, "prompt") # button - self.button = components.icon_button(self, 0, 11, "...", lambda: open_command(self.i, self.ui_state)) - self.button.configure(width=40) + self.button = self.components.icon_button(frame, 0, 11, "...", lambda: open_command(self.i, ui_state)) - self.__set_enabled() + self._bind_save(save_command) + self._set_enabled() - def __switch_enabled(self): + @abstractmethod + def _bind_save(self, save_command): pass + + # BaseConfigListView calls configure_element() on all widget types generically; + # sampling widgets have no post-window logic, so this is an intentional no-op. + def configure_element(self): pass # noqa: B027 + + def _switch_enabled(self): self.save_command() - self.__set_enabled() + self._set_enabled() - def __set_enabled(self): + def _set_enabled(self): enabled = self.element.enabled self.width_entry.configure(state="normal" if enabled else "disabled") self.height_entry.configure(state="normal" if enabled else "disabled") self.prompt_entry.configure(state="normal" if enabled else "disabled") self.seed_entry.configure(state="normal" if enabled else "disabled") self.button.configure(state="normal" if enabled else "disabled") - - def configure_element(self): - pass - - def place_in_list(self): - self.grid(row=self.i, column=0, pady=5, padx=5, sticky="new") diff --git a/modules/ui/BaseSchedulerParamsWindowView.py b/modules/ui/BaseSchedulerParamsWindowView.py index f96ed4876..1106d5227 100644 --- a/modules/ui/BaseSchedulerParamsWindowView.py +++ b/modules/ui/BaseSchedulerParamsWindowView.py @@ -1,119 +1,21 @@ -from modules.ui.ConfigList import ConfigList -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.LearningRateScheduler import LearningRateScheduler -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState -import customtkinter as ctk +from modules.ui.BaseConfigListView import BaseConfigListView -class KvParams(ConfigList): - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__( - master, - train_config, - ui_state, - attr_name="scheduler_params", - from_external_file=False, - add_button_text="add parameter", - is_full_width=True - ) +class BaseSchedulerParamsWindowView: + def __init__(self, components): + self.components = components - def refresh_ui(self): - self._create_element_list() + def build_content(self, master, controller, ui_state): + if controller.is_custom_scheduler(): + self.components.label(master, 0, 0, "Class Name", + tooltip="Python class module and name for the custom scheduler class, in the form of ..") + self.components.entry(master, 0, 1, ui_state, "custom_learning_rate_scheduler") - def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - return KvWidget(master, element, i, open_command, remove_command, clone_command, save_command) - def create_new_element(self) -> dict[str, str]: - return {"key": "", "value": ""} +class BaseKvParamsView(BaseConfigListView): + def __init__(self, components): + self.components = components def open_element_window(self, i, ui_state): pass - - -class KvWidget(ctk.CTkFrame): - def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): - super().__init__(master=master, bg_color="transparent") - self.element = element - self.ui_state = UIState(self, element) - self.i = i - self.save_command = save_command - - self.grid_columnconfigure(0, weight=0) - self.grid_columnconfigure(1, weight=1, uniform=1) - self.grid_columnconfigure(2, weight=1, uniform=1) - - close_button = ctk.CTkButton( - master=self, - width=20, - height=20, - text="X", - corner_radius=2, - fg_color="#C00000", - command=lambda: remove_command(self.i)) - close_button.grid(row=0, column=0) - - # Key - tooltip_key = "Key name for an argument in your scheduler" - self.key = components.entry(self, 0, 1, self.ui_state, "key", - tooltip=tooltip_key, wide_tooltip=True) - self.key.bind("", lambda _: save_command()) - self.key.configure(width=50) - - # Value - tooltip_val = "Value for an argument in your scheduler. Some special values can be used, wrapped in percent signs: LR, EPOCHS, STEPS_PER_EPOCH, TOTAL_STEPS, SCHEDULER_STEPS. Note that OneTrainer calls step() after every individual learning step, not every epoch, so what Torch calls 'epoch' you should treat as 'step'." - self.value = components.entry(self, 0, 2, self.ui_state, "value", - tooltip=tooltip_val, wide_tooltip=True) - self.value.bind("", lambda _: save_command()) - self.value.configure(width=50) - - def place_in_list(self): - self.grid(row=self.i, column=0, padx=5, pady=5, sticky="new") - - -class SchedulerParamsWindow(ctk.CTkToplevel): - def __init__(self, parent, train_config: TrainConfig, ui_state, *args, **kwargs): - super().__init__(parent, *args, **kwargs) - - self.parent = parent - self.train_config = train_config - self.ui_state = ui_state - - self.title("Learning Rate Scheduler Settings") - self.geometry("800x400") - self.resizable(True, True) - - self.grid_rowconfigure(0, weight=1) - self.grid_rowconfigure(1, weight=0) - self.grid_columnconfigure(0, weight=1) - - self.frame = ctk.CTkFrame(self) - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) - - self.expand_frame = ctk.CTkFrame(self.frame, bg_color="transparent") - self.expand_frame.grid(row=1, column=0, columnspan=2, sticky="nsew") - - components.button(self, 1, 0, "ok", command=self.on_window_close) - self.main_frame(self.frame) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def main_frame(self, master): - if self.train_config.learning_rate_scheduler is LearningRateScheduler.CUSTOM: - components.label(master, 0, 0, "Class Name", - tooltip="Python class module and name for the custom scheduler class, in the form of ..") - components.entry(master, 0, 1, self.ui_state, "custom_learning_rate_scheduler") - - # Any additional parameters, in key-value form. - self.params = KvParams(self.expand_frame, self.train_config, self.ui_state) - - def on_window_close(self): - self.destroy() diff --git a/modules/ui/BaseTimestepDistributionWindowView.py b/modules/ui/BaseTimestepDistributionWindowView.py index 21e41ce3e..19354bf43 100644 --- a/modules/ui/BaseTimestepDistributionWindowView.py +++ b/modules/ui/BaseTimestepDistributionWindowView.py @@ -1,186 +1,46 @@ -from modules.modelSetup.mixin.ModelSetupNoiseMixin import ( - ModelSetupNoiseMixin, -) -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.TimestepDistribution import TimestepDistribution -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState -import torch -from torch import Tensor -import customtkinter as ctk -from customtkinter import AppearanceModeTracker, ThemeManager -from matplotlib import pyplot as plt -from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg -class TimestepGenerator(ModelSetupNoiseMixin): - - def __init__( - self, - timestep_distribution: TimestepDistribution, - min_noising_strength: float, - max_noising_strength: float, - noising_weight: float, - noising_bias: float, - timestep_shift: float, - ): - super().__init__() - - self.timestep_distribution = timestep_distribution - self.min_noising_strength = min_noising_strength - self.max_noising_strength = max_noising_strength - self.noising_weight = noising_weight - self.noising_bias = noising_bias - self.timestep_shift = timestep_shift - - def generate(self) -> Tensor: - generator = torch.Generator() - generator.seed() - - config = TrainConfig.default_values() - config.timestep_distribution = self.timestep_distribution - config.min_noising_strength = self.min_noising_strength - config.max_noising_strength = self.max_noising_strength - config.noising_weight = self.noising_weight - config.noising_bias = self.noising_bias - config.timestep_shift = self.timestep_shift - - - return self._get_timestep_discrete( - num_train_timesteps=1000, - deterministic=False, - generator=generator, - batch_size=1000000, - config=config, - ) - - -class TimestepDistributionWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - config: TrainConfig, - ui_state: UIState, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) - - self.title("Timestep Distribution") - self.geometry("900x600") - self.resizable(True, True) - - self.config = config - self.ui_state = ui_state - self.image_preview_file_index = 0 - self.ax = None - self.canvas = None - - self.grid_rowconfigure(0, weight=1) - self.grid_columnconfigure(0, weight=1) - - frame = self.__content_frame(self) - frame.grid(row=0, column=0, sticky='nsew') - components.button(self, 1, 0, "ok", self.__ok) - - self.wait_visibility() - self.after(200, lambda: set_window_icon(self)) - self.grab_set() - self.focus_set() - - def __content_frame(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - frame.grid_rowconfigure(7, weight=1) +class BaseTimestepDistributionWindowView: + def __init__(self, components): + self.components = components + def build_content(self, frame, controller, ui_state): # timestep distribution - components.label(frame, 0, 0, "Timestep Distribution", + self.components.label(frame, 0, 0, "Timestep Distribution", tooltip="Selects the function to sample timesteps during training", wide_tooltip=True) - components.options(frame, 0, 1, [str(x) for x in list(TimestepDistribution)], self.ui_state, + self.components.options(frame, 0, 1, controller.get_distribution_options(), ui_state, "timestep_distribution") # min noising strength - components.label(frame, 1, 0, "Min Noising Strength", + self.components.label(frame, 1, 0, "Min Noising Strength", tooltip="Specifies the minimum noising strength used during training. This can help to improve composition, but prevents finer details from being trained") - components.entry(frame, 1, 1, self.ui_state, "min_noising_strength") + self.components.entry(frame, 1, 1, ui_state, "min_noising_strength") # max noising strength - components.label(frame, 2, 0, "Max Noising Strength", + self.components.label(frame, 2, 0, "Max Noising Strength", tooltip="Specifies the maximum noising strength used during training. This can be useful to reduce overfitting, but also reduces the impact of training samples on the overall image composition") - components.entry(frame, 2, 1, self.ui_state, "max_noising_strength") + self.components.entry(frame, 2, 1, ui_state, "max_noising_strength") # noising weight - components.label(frame, 3, 0, "Noising Weight", + self.components.label(frame, 3, 0, "Noising Weight", tooltip="Controls the weight parameter of the timestep distribution function. Use the preview to see more details.") - components.entry(frame, 3, 1, self.ui_state, "noising_weight") + self.components.entry(frame, 3, 1, ui_state, "noising_weight") # noising bias - components.label(frame, 4, 0, "Noising Bias", + self.components.label(frame, 4, 0, "Noising Bias", tooltip="Controls the bias parameter of the timestep distribution function. Use the preview to see more details.") - components.entry(frame, 4, 1, self.ui_state, "noising_bias") + self.components.entry(frame, 4, 1, ui_state, "noising_bias") # timestep shift - components.label(frame, 5, 0, "Timestep Shift", + self.components.label(frame, 5, 0, "Timestep Shift", tooltip="Shift the timestep distribution. Use the preview to see more details.") - components.entry(frame, 5, 1, self.ui_state, "timestep_shift") + self.components.entry(frame, 5, 1, ui_state, "timestep_shift") # dynamic timestep shifting - components.label(frame, 6, 0, "Dynamic Timestep Shifting", + self.components.label(frame, 6, 0, "Dynamic Timestep Shifting", tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Dynamic Timestep Shifting is not shown in the preview. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) - components.switch(frame, 6, 1, self.ui_state, "dynamic_timestep_shifting") - - - # plot - appearance_mode = AppearanceModeTracker.get_mode() - background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) - text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) - background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" - text_color = f"#{int(text_color[0]/256):x}{int(text_color[1]/256):x}{int(text_color[2]/256):x}" - - fig, ax = plt.subplots() - self.ax = ax - self.canvas = FigureCanvasTkAgg(fig, master=frame) - self.canvas.get_tk_widget().grid(row=0, column=3, rowspan=8) - - fig.set_facecolor(background_color) - ax.set_facecolor(background_color) - ax.spines['bottom'].set_color(text_color) - ax.spines['left'].set_color(text_color) - ax.spines['top'].set_color(text_color) - ax.spines['right'].set_color(text_color) - ax.tick_params(axis='x', colors=text_color, which="both") - ax.tick_params(axis='y', colors=text_color, which="both") - ax.xaxis.label.set_color(text_color) - ax.yaxis.label.set_color(text_color) - - self.__update_preview() - - # update button - components.button(frame, 8, 3, "Update Preview", command=self.__update_preview) - - frame.pack(fill="both", expand=1) - return frame - - def __update_preview(self): - generator = TimestepGenerator( - timestep_distribution=self.config.timestep_distribution, - min_noising_strength=self.config.min_noising_strength, - max_noising_strength=self.config.max_noising_strength, - noising_weight=self.config.noising_weight, - noising_bias=self.config.noising_bias, - timestep_shift=self.config.timestep_shift, - ) - - self.ax.cla() - self.ax.hist(generator.generate(), bins=1000, range=(0, 999)) - self.canvas.draw() - - def __ok(self): - self.destroy() + self.components.switch(frame, 6, 1, ui_state, "dynamic_timestep_shifting") diff --git a/modules/ui/BaseTopBarView.py b/modules/ui/BaseTopBarView.py index 820fdb71a..f50c1ce14 100644 --- a/modules/ui/BaseTopBarView.py +++ b/modules/ui/BaseTopBarView.py @@ -1,7 +1,7 @@ import json import os import traceback -import webbrowser +from abc import abstractmethod from collections.abc import Callable from contextlib import suppress @@ -11,25 +11,45 @@ from modules.util.enum.ModelType import ModelType from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.optimizer_util import change_optimizer -from modules.util.path_util import write_json_atomic -from modules.util.ui import components, dialogs -from modules.util.ui.UIState import UIState -import customtkinter as ctk +class BaseTopBarView: + def __init__(self, components): + self.components = components -class TopBar: - def __init__( + @abstractmethod + def _make_config_ui_state(self, master, data): + pass + + @abstractmethod + def _get_dropdown_text(self, widget) -> str: + pass + + @abstractmethod + def _setup_frame_column_weight(self): + pass + + @abstractmethod + def _forget_dropdown(self): + pass + + @abstractmethod + def _show_save_dialog(self, default_value: str, callback): + pass + + def build( self, + frame, master, - train_config: TrainConfig, - ui_state: UIState, + controller, + ui_state, change_model_type_callback: Callable[[ModelType], None], change_training_method_callback: Callable[[TrainingMethod], None], load_preset_callback: Callable[[], None], ): + self.controller = controller + self.frame = frame self.master = master - self.train_config = train_config self.ui_state = ui_state self.change_model_type_callback = change_model_type_callback self.change_training_method_callback = change_training_method_callback @@ -40,20 +60,16 @@ def __init__( self.config_ui_data = { "config_name": path_util.canonical_join(self.dir, "#.json") } - self.config_ui_state = UIState(master, self.config_ui_data) + self.config_ui_state = self._make_config_ui_state(master, self.config_ui_data) - self.configs = [("", path_util.canonical_join(self.dir, "#.json"))] - self.__load_available_config_names() + self.configs = controller.load_available_config_names(self.dir) self.current_config = [] - self.frame = ctk.CTkFrame(master=master, corner_radius=0) - self.frame.grid(row=0, column=0, sticky="nsew") - self.training_method = None # title - components.app_title(self.frame, 0, 0) + self.components.app_title(self.frame, 0, 0) # dropdown self.configs_dropdown = None @@ -61,49 +77,25 @@ def __init__( # remove button # TODO - # components.icon_button(self.frame, 0, 2, "-", self.__remove_config) + # self.components.icon_button(self.frame, 0, 2, "-", self.__remove_config) # Wiki button - components.button(self.frame, 0, 4, "Wiki", self.open_wiki, width=50) + self.components.button(self.frame, 0, 4, "Wiki", self.open_wiki, width=50) # save button - components.button(self.frame, 0, 3, "Save config", self.__save_config, - tooltip="Save the current configuration in a custom preset", width=90) + self.components.button(self.frame, 0, 3, "Save config", self.__save_config, + tooltip="Save the current configuration in a custom preset", width=90) # padding - self.frame.grid_columnconfigure(5, weight=1) + self._setup_frame_column_weight() # model type - components.options_kv( + self.components.options_kv( master=self.frame, row=0, column=6, - values=[ #TODO simplify - ("SD1.5", ModelType.STABLE_DIFFUSION_15), - ("SD1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), - ("SD2.0", ModelType.STABLE_DIFFUSION_20), - ("SD2.0 Inpainting", ModelType.STABLE_DIFFUSION_20_INPAINTING), - ("SD2.1", ModelType.STABLE_DIFFUSION_21), - ("SD3", ModelType.STABLE_DIFFUSION_3), - ("SD3.5", ModelType.STABLE_DIFFUSION_35), - ("SDXL", ModelType.STABLE_DIFFUSION_XL_10_BASE), - ("SDXL Inpainting", ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING), - ("Wuerstchen v2", ModelType.WUERSTCHEN_2), - ("Stable Cascade", ModelType.STABLE_CASCADE_1), - ("PixArt Alpha", ModelType.PIXART_ALPHA), - ("PixArt Sigma", ModelType.PIXART_SIGMA), - ("Flux Dev.1", ModelType.FLUX_DEV_1), - ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), - ("Flux 2 [Dev, Klein]", ModelType.FLUX_2), - ("Sana", ModelType.SANA), - ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), - ("HiDream Full", ModelType.HI_DREAM_FULL), - ("Chroma1", ModelType.CHROMA_1), - ("QwenImage", ModelType.QWEN), - ("Z-Image", ModelType.Z_IMAGE), - ("Ernie Image", ModelType.ERNIE), - ], - ui_state=self.ui_state, + values=controller.get_model_types(), + ui_state=ui_state, var_name="model_type", command=self.__change_model_type, ) @@ -112,40 +104,9 @@ def __create_training_method(self): if self.training_method: self.training_method.destroy() - values = [] - #TODO simplify - if self.train_config.model_type.is_stable_diffusion(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ("Embedding", TrainingMethod.EMBEDDING), - ("Fine Tune VAE", TrainingMethod.FINE_TUNE_VAE), - ] - elif self.train_config.model_type.is_stable_diffusion_3() \ - or self.train_config.model_type.is_stable_diffusion_xl() \ - or self.train_config.model_type.is_wuerstchen() \ - or self.train_config.model_type.is_pixart() \ - or self.train_config.model_type.is_flux_1() \ - or self.train_config.model_type.is_sana() \ - or self.train_config.model_type.is_hunyuan_video() \ - or self.train_config.model_type.is_hi_dream() \ - or self.train_config.model_type.is_chroma(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ("Embedding", TrainingMethod.EMBEDDING), - ] - elif self.train_config.model_type.is_qwen() \ - or self.train_config.model_type.is_z_image() \ - or self.train_config.model_type.is_flux_2() \ - or self.train_config.model_type.is_ernie(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ] - - # training method - self.training_method = components.options_kv( + values = self.controller.get_training_methods(self.controller.train_config.model_type) + + self.training_method = self.components.options_kv( master=self.frame, row=0, column=7, @@ -161,40 +122,21 @@ def __change_model_type(self, model_type: ModelType): def __create_configs_dropdown(self): if self.configs_dropdown is not None: - self.configs_dropdown.grid_forget() + self._forget_dropdown() - self.configs_dropdown = components.options_kv( + self.configs_dropdown = self.components.options_kv( self.frame, 0, 1, self.configs, self.config_ui_state, "config_name", self.__load_current_config ) - def __load_available_config_names(self): - if os.path.isdir(self.dir): - for path in os.listdir(self.dir): - if path != "#.json": - path = path_util.canonical_join(self.dir, path) - if path.endswith(".json") and os.path.isfile(path): - name = os.path.basename(path) - name = os.path.splitext(name)[0] - self.configs.append((name, path)) - self.configs.sort() - - def __save_to_file(self, name) -> str: - name = path_util.safe_filename(name) - path = path_util.canonical_join("training_presets", f"{name}.json") - - write_json_atomic(path, self.train_config.to_settings_dict(secrets=False)) - - return path - - def __save_secrets(self, path) -> str: - write_json_atomic(path, self.train_config.secrets.to_dict()) - return path + def __save_config(self): + default_value = self._get_dropdown_text(self.configs_dropdown) + while default_value.startswith('#'): + default_value = default_value[1:] - def open_wiki(self): - webbrowser.open("https://github.com/Nerogar/OneTrainer/wiki", new=0, autoraise=False) + self._show_save_dialog(default_value, self.__save_new_config) def __save_new_config(self, name): - path = self.__save_to_file(name) + path = self.controller.save_to_file(name) is_new_config = name not in [x[0] for x in self.configs] @@ -208,20 +150,6 @@ def __save_new_config(self, name): if is_new_config: self.__create_configs_dropdown() - def __save_config(self): - default_value = self.configs_dropdown.get() - while default_value.startswith('#'): - default_value = default_value[1:] - - dialogs.StringInputDialog( - parent=self.master, - title="name", - question="Config Name", - callback=self.__save_new_config, - default_value=default_value, - validate_callback=lambda x: not x.startswith("#") - ) - def __load_current_config(self, filename): try: basename = os.path.basename(filename) @@ -239,10 +167,10 @@ def __load_current_config(self, filename): secrets_dict=json.load(f) loaded_config.secrets = SecretsConfig.default_values().from_dict(secrets_dict) - self.train_config.from_dict(loaded_config.to_dict()) + self.controller.train_config.from_dict(loaded_config.to_dict()) self.ui_state.update(loaded_config) - optimizer_config = change_optimizer(self.train_config) + optimizer_config = change_optimizer(self.controller.train_config) self.ui_state.get_var("optimizer").update(optimizer_config) self.load_preset_callback() @@ -255,6 +183,8 @@ def __remove_config(self): # TODO pass + def open_wiki(self): + self.controller.open_wiki() + def save_default(self): - self.__save_to_file("#") - self.__save_secrets("secrets.json") + self.controller.save_default() diff --git a/modules/ui/BaseTrainUIView.py b/modules/ui/BaseTrainUIView.py index ba90d2e64..7acbf0005 100644 --- a/modules/ui/BaseTrainUIView.py +++ b/modules/ui/BaseTrainUIView.py @@ -1,889 +1,356 @@ -import ctypes -import datetime -import json -import os -import platform -import subprocess -import sys -import threading -import time -import traceback -import webbrowser +from abc import ABC, abstractmethod from collections.abc import Callable -from contextlib import suppress -from pathlib import Path -from tkinter import filedialog, messagebox - -import scripts.generate_debug_report -from modules.ui.AdditionalEmbeddingsTab import AdditionalEmbeddingsTab -from modules.ui.CaptionUI import CaptionUI -from modules.ui.CloudTab import CloudTab -from modules.ui.ConceptTab import ConceptTab -from modules.ui.ConvertModelUI import ConvertModelUI -from modules.ui.LoraTab import LoraTab -from modules.ui.ModelTab import ModelTab -from modules.ui.ProfilingWindow import ProfilingWindow -from modules.ui.SampleWindow import SampleWindow -from modules.ui.SamplingTab import SamplingTab -from modules.ui.TopBar import TopBar -from modules.ui.TrainingTab import TrainingTab -from modules.ui.VideoToolUI import VideoToolUI -from modules.util import create -from modules.util.callbacks.TrainCallbacks import TrainCallbacks -from modules.util.commands.TrainCommands import TrainCommands -from modules.util.config.TrainConfig import TrainConfig + +from modules.util import path_util from modules.util.enum.DataType import DataType 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.PathIOType import PathIOType -from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.torch_util import torch_gc -from modules.util.TrainProgress import TrainProgress -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState -from modules.util.ui.validation import flush_and_validate_all - -import torch - -import customtkinter as ctk -from customtkinter import AppearanceModeTracker - -# 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": - with suppress(Exception): - # 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): - set_step_progress: Callable[[int, int], None] - set_epoch_progress: Callable[[int, int], None] - - status_label: ctk.CTkLabel | None - training_button: ctk.CTkButton | None - training_callbacks: TrainCallbacks | None - training_commands: TrainCommands | None - - _TRAIN_BUTTON_STYLES = { - "idle": { - "text": "Start Training", - "state": "normal", - "fg_color": "#198754", - "hover_color": "#146c43", - "text_color": "white", - "text_color_disabled": "white", - }, - "running": { - "text": "Stop Training", - "state": "normal", - "fg_color": "#dc3545", - "hover_color": "#bb2d3b", - "text_color": "white", - }, - "stopping": { - "text": "Stopping...", - "state": "disabled", - "fg_color": "#dc3545", - "hover_color": "#dc3545", - "text_color": "white", - "text_color_disabled": "white", - }, - } - - def __init__(self): - super().__init__() - - self.title("OneTrainer") - self.geometry("1100x740") - - self.after(100, lambda: self._set_icon()) - - # more efficient version of ctk.set_appearance_mode("System"), which retrieves the system theme on each main loop iteration - ctk.set_appearance_mode("Light" if AppearanceModeTracker.detect_appearance_mode() == 0 else "Dark") - ctk.set_default_color_theme("blue") - - self.train_config = TrainConfig.default_values() - self.ui_state = UIState(self, self.train_config) - - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=1) - self.grid_rowconfigure(2, weight=0) - self.grid_columnconfigure(0, weight=1) - - self.status_label = None - self.eta_label = None - self.training_button = None - self.export_button = None - self.tabview = None - - self.model_tab = None - self.training_tab = None - self.lora_tab = None - self.cloud_tab = None - self.additional_embeddings_tab = None - - self.top_bar_component = self.top_bar(self) - self.content_frame(self) - self.bottom_bar(self) - - 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.workspace_dir_trace_id = self.ui_state.add_var_trace("workspace_dir", self._on_workspace_dir_change_trace) - - # Persistent profiling window. - self.profiling_window = ProfilingWindow(self) - - self.protocol("WM_DELETE_WINDOW", self.__close) - - def __close(self): - self.top_bar_component.save_default() - self._stop_always_on_tensorboard() - if hasattr(self, 'workspace_dir_trace_id'): - self.ui_state.remove_var_trace("workspace_dir", self.workspace_dir_trace_id) - self.quit() - - def top_bar(self, master): - return TopBar( - master, - self.train_config, - self.ui_state, - self.change_model_type, - self.change_training_method, - self.load_preset, - ) - def _set_icon(self): - """Set the window icon safely after window is ready""" - set_window_icon(self) - def bottom_bar(self, master): - frame = ctk.CTkFrame(master=master, corner_radius=0) - frame.grid(row=2, column=0, sticky="nsew") +class BaseTrainUIView(ABC): + def __init__(self, components): + self.components = components - self.set_step_progress, self.set_epoch_progress = components.double_progress(frame, 0, 0, "step", "epoch") + # --- Abstract callbacks (controller calls into view) --- - # status + ETA container - self.status_frame = ctk.CTkFrame(frame, corner_radius=0, fg_color="transparent") - self.status_frame.grid(row=0, column=1, sticky="w") - self.status_frame.grid_rowconfigure(0, weight=0) - self.status_frame.grid_rowconfigure(1, weight=0) - self.status_frame.grid_columnconfigure(0, weight=1) + @abstractmethod + def on_update_status(self, status: str): pass - self.status_label = components.label(self.status_frame, 0, 0, "", pad=0, - tooltip="Current status of the training run") - self.eta_label = components.label(self.status_frame, 1, 0, "", pad=0) + @abstractmethod + def on_training_started(self): pass - # padding - frame.grid_columnconfigure(2, weight=1) + @abstractmethod + def on_training_stopped(self, error_caught: bool): pass + @abstractmethod + def on_training_stopping(self): pass - # export button - self.export_button = components.button(frame, 0, 3, "Export", self.export_training, - width=60, padx=5, pady=(15, 0), - tooltip="Export the current configuration as a script to run without a UI") + @abstractmethod + def on_update_progress(self, epoch_step: int, max_step: int, epoch: int, max_epoch: int, eta_str: str | None): pass - # debug button - components.button(frame, 0, 4, "Debug", self.generate_debug_package, - width=60, padx=(5, 25), pady=(15, 0), - tooltip="Generate a zip file with config.json, debug_report.log and settings diff, use this to report bugs or issues") + @abstractmethod + def schedule_on_main_thread(self, fn: Callable): pass - # tensorboard button - components.button(frame, 0, 5, "Tensorboard", self.open_tensorboard, - width=100, padx=(0, 5), pady=(15, 0)) + @abstractmethod + def get_cloud_reattach(self) -> bool: pass - # training button - self.training_button = components.button(frame, 0, 6, "Start Training", self.start_training, - padx=(5, 20), pady=(15, 0)) - self._set_training_button_style("idle") # centralized styling + @abstractmethod + def save_default(self): pass + + @abstractmethod + def show_validation_errors(self, errors: list[str]): pass + + @abstractmethod + def wait_window(self, window): pass + + @abstractmethod + def show_window(self, window): pass + + @abstractmethod + def connect_window_closed(self, window, callback): pass + + def sync_cloud_secrets(self): + self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) - return frame + def start_training(self): + self.controller.start_training() + + def open_tensorboard(self): + self.controller.open_tensorboard() - def content_frame(self, master): - frame = ctk.CTkFrame(master=master, corner_radius=0) - frame.grid(row=1, column=0, sticky="nsew") + def sample_now(self): + self.controller.sample_now() - frame.grid_rowconfigure(0, weight=1) - frame.grid_columnconfigure(0, weight=1) + def backup_now(self): + self.controller.backup_now() - self.tabview = ctk.CTkTabview(frame) - self.tabview.grid(row=0, column=0, sticky="nsew") + def save_now(self): + self.controller.save_now() - 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")) + @abstractmethod + def open_dataset_tool(self): pass - self.change_training_method(self.train_config.training_method) + @abstractmethod + def open_video_tool(self): pass - return frame + @abstractmethod + def open_convert_model_tool(self): pass - def create_general_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) + @abstractmethod + def open_sampling_tool(self): pass + + @abstractmethod + def open_manual_sample_window(self): pass + + # --- Content builders (components calls; called by CTK view after frame creation) --- + + def build_bottom_bar_content(self, frame, status_frame, controller, ui_state): + self.set_step_progress, self.set_epoch_progress = self.components.double_progress(frame, 0, 0, "step", "epoch") + + self.status_label = self.components.label(status_frame, 0, 0, "", pad=0, + tooltip="Current status of the training run") + self.eta_label = self.components.label(status_frame, 1, 0, "", pad=0) + self.export_button = self.components.button(frame, 0, 3, "Export", self.export_training, + width=60, padx=5, pady=(15, 0), + tooltip="Export the current configuration as a script to run without a UI") + + self.components.button(frame, 0, 4, "Debug", self.generate_debug_package, + width=60, padx=(5, 25), pady=(15, 0), + tooltip="Generate a zip file with config.json, debug_report.log and settings diff, use this to report bugs or issues") + + self.components.button(frame, 0, 5, "Tensorboard", self.open_tensorboard, + width=100, padx=(0, 5), pady=(15, 0)) + + self.training_button = self.components.button(frame, 0, 6, "Start Training", self.start_training, + padx=(5, 20), pady=(15, 0)) + + def build_general_tab_content(self, frame, controller, ui_state): # workspace dir - components.label(frame, 0, 0, "Workspace Directory", + self.components.label(frame, 0, 0, "Workspace Directory", tooltip="The directory where all files of this training run are saved") - components.path_entry(frame, 0, 1, self.ui_state, "workspace_dir", mode="dir", command=self._on_workspace_dir_change) + self.components.path_entry(frame, 0, 1, ui_state, "workspace_dir", mode="dir", command=controller._on_workspace_dir_change) # cache dir - components.label(frame, 0, 2, "Cache Directory", + self.components.label(frame, 0, 2, "Cache Directory", tooltip="The directory where cached data is saved") - components.path_entry(frame, 0, 3, self.ui_state, "cache_dir", mode="dir") + self.components.path_entry(frame, 0, 3, ui_state, "cache_dir", mode="dir") # continue from previous backup - components.label(frame, 2, 0, "Continue from last backup", + self.components.label(frame, 2, 0, "Continue from last backup", tooltip="Automatically continues training from the last backup saved in /backup") - components.switch(frame, 2, 1, self.ui_state, "continue_last_backup") + self.components.switch(frame, 2, 1, ui_state, "continue_last_backup") # only cache - components.label(frame, 2, 2, "Only Cache", + self.components.label(frame, 2, 2, "Only Cache", tooltip="Only populate the cache, without any training") - components.switch(frame, 2, 3, self.ui_state, "only_cache") + self.components.switch(frame, 2, 3, ui_state, "only_cache") # TODO: In Phase 4 rework the general tab. # prevent overwrites - components.label(frame, 3, 0, "Prevent Overwrites", + self.components.label(frame, 3, 0, "Prevent Overwrites", tooltip="When enabled, output paths that already exist on disk will be flagged as invalid to avoid accidental overwrites") - components.switch(frame, 3, 1, self.ui_state, "prevent_overwrites") + self.components.switch(frame, 3, 1, ui_state, "prevent_overwrites") # debug - components.label(frame, 4, 0, "Debug mode", + self.components.label(frame, 4, 0, "Debug mode", tooltip="Save debug information during the training into the debug directory") - components.switch(frame, 4, 1, self.ui_state, "debug_mode") + self.components.switch(frame, 4, 1, ui_state, "debug_mode") - components.label(frame, 4, 2, "Debug Directory", + self.components.label(frame, 4, 2, "Debug Directory", tooltip="The directory where debug data is saved") - components.path_entry(frame, 4, 3, self.ui_state, "debug_dir", mode="dir", io_type=PathIOType.OUTPUT) + self.components.path_entry(frame, 4, 3, ui_state, "debug_dir", mode="dir", io_type=PathIOType.OUTPUT) # tensorboard - components.label(frame, 6, 0, "Tensorboard", + self.components.label(frame, 6, 0, "Tensorboard", tooltip="Starts the Tensorboard Web UI during training") - components.switch(frame, 6, 1, self.ui_state, "tensorboard") + self.components.switch(frame, 6, 1, ui_state, "tensorboard") - components.label(frame, 6, 2, "Always-On Tensorboard", + self.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) + self.components.switch(frame, 6, 3, ui_state, "tensorboard_always_on", command=controller._on_always_on_tensorboard_toggle) - components.label(frame, 7, 0, "Expose Tensorboard", + self.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", + self.components.switch(frame, 7, 1, ui_state, "tensorboard_expose") + self.components.label(frame, 7, 2, "Tensorboard Port", tooltip="Port to use for Tensorboard link") - components.entry(frame, 7, 3, self.ui_state, "tensorboard_port") - + self.components.entry(frame, 7, 3, ui_state, "tensorboard_port") # validation - components.label(frame, 8, 0, "Validation", + self.components.label(frame, 8, 0, "Validation", tooltip="Enable validation steps and add new graph in tensorboard") - components.switch(frame, 8, 1, self.ui_state, "validation") + self.components.switch(frame, 8, 1, ui_state, "validation") - components.label(frame, 8, 2, "Validate after", + self.components.label(frame, 8, 2, "Validate after", tooltip="The interval used when validate training") - components.time_entry(frame, 8, 3, self.ui_state, "validate_after", "validate_after_unit") + self.components.time_entry(frame, 8, 3, ui_state, "validate_after", "validate_after_unit") # device - components.label(frame, 10, 0, "Dataloader Threads", + self.components.label(frame, 10, 0, "Dataloader Threads", tooltip="Number of threads used for the data loader. Increase if your GPU has room during caching, decrease if it's going out of memory during caching.") - components.entry(frame, 10, 1, self.ui_state, "dataloader_threads", required=True) + self.components.entry(frame, 10, 1, ui_state, "dataloader_threads", required=True) - components.label(frame, 11, 0, "Train Device", + self.components.label(frame, 11, 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", required=True) + self.components.entry(frame, 11, 1, ui_state, "train_device", required=True) - components.label(frame, 12, 0, "Multi-GPU", + self.components.label(frame, 12, 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", + self.components.switch(frame, 12, 1, ui_state, "multi_gpu") + self.components.label(frame, 12, 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") + self.components.entry(frame, 12, 3, ui_state, "device_indexes") - components.label(frame, 13, 0, "Gradient Reduce Precision", + self.components.label(frame, 13, 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, + self.components.options(frame, 13, 1, [str(x) for x in list(GradientReducePrecision)], ui_state, "gradient_reduce_precision") - components.label(frame, 13, 2, "Fused Gradient Reduce", + self.components.label(frame, 13, 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") + self.components.switch(frame, 13, 3, ui_state, "fused_gradient_reduce") - components.label(frame, 14, 0, "Async Gradient Reduce", + self.components.label(frame, 14, 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)", + self.components.switch(frame, 14, 1, ui_state, "async_gradient_reduce") + self.components.label(frame, 14, 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") + self.components.entry(frame, 14, 3, ui_state, "async_gradient_reduce_buffer") - components.label(frame, 15, 0, "Temp Device", + self.components.label(frame, 15, 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") - - frame.pack(fill="both", expand=1) - return frame - - def create_model_tab(self, master): - return ModelTab(master, self.train_config, self.ui_state) - - def create_data_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) + self.components.entry(frame, 15, 1, ui_state, "temp_device") + def build_data_tab_content(self, frame, controller, ui_state): # aspect ratio bucketing - components.label(frame, 0, 0, "Aspect Ratio Bucketing", + self.components.label(frame, 0, 0, "Aspect Ratio Bucketing", tooltip="Aspect ratio bucketing enables training on images with different aspect ratios") - components.switch(frame, 0, 1, self.ui_state, "aspect_ratio_bucketing") + self.components.switch(frame, 0, 1, ui_state, "aspect_ratio_bucketing") # latent caching - components.label(frame, 1, 0, "Latent Caching", + self.components.label(frame, 1, 0, "Latent Caching", tooltip="Caching of intermediate training data that can be re-used between epochs") - components.switch(frame, 1, 1, self.ui_state, "latent_caching") + self.components.switch(frame, 1, 1, ui_state, "latent_caching") # clear cache before training - components.label(frame, 2, 0, "Clear cache before training", + self.components.label(frame, 2, 0, "Clear cache before training", tooltip="Clears the cache directory before starting to train. Only disable this if you want to continue using the same cached data. Disabling this can lead to errors, if other settings are changed during a restart") - components.switch(frame, 2, 1, self.ui_state, "clear_cache_before_training") + self.components.switch(frame, 2, 1, ui_state, "clear_cache_before_training") - frame.pack(fill="both", expand=1) - return frame - - def create_concepts_tab(self, master): - return ConceptTab(master, self.train_config, self.ui_state) - - def create_training_tab(self, master) -> TrainingTab: - return TrainingTab(master, self.train_config, self.ui_state) - - def create_cloud_tab(self, master) -> CloudTab: - return CloudTab(master, self.train_config, self.ui_state,parent=self) - - def create_sampling_tab(self, master): - master.grid_rowconfigure(0, weight=0) - master.grid_rowconfigure(1, weight=1) - master.grid_columnconfigure(0, weight=1) - - # sample after - top_frame = ctk.CTkFrame(master=master, corner_radius=0) - top_frame.grid(row=0, column=0, sticky="nsew") - sub_frame = ctk.CTkFrame(master=top_frame, corner_radius=0, fg_color="transparent") - sub_frame.grid(row=1, column=0, sticky="nsew", columnspan=6) - - components.label(top_frame, 0, 0, "Sample After", + def build_sampling_tab_header(self, top_frame, sub_frame, controller, ui_state): + self.components.label(top_frame, 0, 0, "Sample After", tooltip="The interval used when automatically sampling from the model during training") - components.time_entry(top_frame, 0, 1, self.ui_state, "sample_after", "sample_after_unit") + self.components.time_entry(top_frame, 0, 1, ui_state, "sample_after", "sample_after_unit") - components.label(top_frame, 0, 2, "Skip First", + self.components.label(top_frame, 0, 2, "Skip First", tooltip="Start sampling automatically after this interval has elapsed.") - components.entry(top_frame, 0, 3, self.ui_state, "sample_skip_first", width=50, sticky="nw") + self.components.entry(top_frame, 0, 3, ui_state, "sample_skip_first", width=50, sticky="nw") - components.label(top_frame, 0, 4, "Format", + self.components.label(top_frame, 0, 4, "Format", tooltip="File Format used when saving samples") - components.options_kv(top_frame, 0, 5, [ + self.components.options_kv(top_frame, 0, 5, [ ("PNG", ImageFormat.PNG), ("JPG", ImageFormat.JPG), - ], self.ui_state, "sample_image_format") + ], ui_state, "sample_image_format") - components.button(top_frame, 0, 6, "sample now", self.sample_now) + self.components.button(top_frame, 0, 6, "sample now", self.sample_now) - components.button(top_frame, 0, 7, "manual sample", self.open_manual_sample_window ) + self.components.button(top_frame, 0, 7, "manual sample", self.open_manual_sample_window) - components.label(sub_frame, 0, 0, "Non-EMA Sampling", + self.components.label(sub_frame, 0, 0, "Non-EMA Sampling", tooltip="Whether to include non-ema sampling when using ema.") - components.switch(sub_frame, 0, 1, self.ui_state, "non_ema_sampling") + self.components.switch(sub_frame, 0, 1, ui_state, "non_ema_sampling") - components.label(sub_frame, 0, 2, "Samples to Tensorboard", + self.components.label(sub_frame, 0, 2, "Samples to Tensorboard", tooltip="Whether to include sample images in the Tensorboard output.") - components.switch(sub_frame, 0, 3, self.ui_state, "samples_to_tensorboard") - - # table - frame = ctk.CTkFrame(master=master, corner_radius=0) - frame.grid(row=1, column=0, sticky="nsew") - - return SamplingTab(frame, self.train_config, self.ui_state) - - def create_backup_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) + self.components.switch(sub_frame, 0, 3, ui_state, "samples_to_tensorboard") + def build_backup_tab_content(self, frame, controller, ui_state): # backup after - components.label(frame, 0, 0, "Backup After", + self.components.label(frame, 0, 0, "Backup After", tooltip="The interval used when automatically creating model backups during training") - components.time_entry(frame, 0, 1, self.ui_state, "backup_after", "backup_after_unit") + self.components.time_entry(frame, 0, 1, ui_state, "backup_after", "backup_after_unit") # backup now - components.button(frame, 0, 3, "backup now", self.backup_now) + self.components.button(frame, 0, 3, "backup now", self.backup_now) # rolling backup - components.label(frame, 1, 0, "Rolling Backup", + self.components.label(frame, 1, 0, "Rolling Backup", tooltip="If rolling backups are enabled, older backups are deleted automatically") - components.switch(frame, 1, 1, self.ui_state, "rolling_backup") + self.components.switch(frame, 1, 1, ui_state, "rolling_backup") # rolling backup count - components.label(frame, 1, 3, "Rolling Backup Count", + self.components.label(frame, 1, 3, "Rolling Backup Count", tooltip="Defines the number of backups to keep if rolling backups are enabled") - components.entry(frame, 1, 4, self.ui_state, "rolling_backup_count") + self.components.entry(frame, 1, 4, ui_state, "rolling_backup_count") # backup before save - components.label(frame, 2, 0, "Backup Before Save", + self.components.label(frame, 2, 0, "Backup Before Save", tooltip="Create a full backup before saving the final model") - components.switch(frame, 2, 1, self.ui_state, "backup_before_save") + self.components.switch(frame, 2, 1, ui_state, "backup_before_save") # save after - components.label(frame, 3, 0, "Save Every", + self.components.label(frame, 3, 0, "Save Every", tooltip="The interval used when automatically saving the model during training") - components.time_entry(frame, 3, 1, self.ui_state, "save_every", "save_every_unit") + self.components.time_entry(frame, 3, 1, ui_state, "save_every", "save_every_unit") # save now - components.button(frame, 3, 3, "save now", self.save_now) + self.components.button(frame, 3, 3, "save now", self.save_now) # skip save - components.label(frame, 4, 0, "Skip First", + self.components.label(frame, 4, 0, "Skip First", tooltip="Start saving automatically after this interval has elapsed") - components.entry(frame, 4, 1, self.ui_state, "save_skip_first", width=50, sticky="nw") + self.components.entry(frame, 4, 1, ui_state, "save_skip_first", width=50, sticky="nw") # save filename prefix - components.label(frame, 5, 0, "Save Filename Prefix", + self.components.label(frame, 5, 0, "Save Filename Prefix", tooltip="The prefix for filenames used when saving the model during training") - components.entry(frame, 5, 1, self.ui_state, "save_filename_prefix") - - frame.pack(fill="both", expand=1) - return frame - - def embedding_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) + self.components.entry(frame, 5, 1, ui_state, "save_filename_prefix") + def build_embedding_tab_content(self, frame, controller, ui_state): # embedding model name - components.label(frame, 0, 0, "Base embedding", + self.components.label(frame, 0, 0, "Base embedding", tooltip="The base embedding to train on. Leave empty to create a new embedding") - components.path_entry( - frame, 0, 1, self.ui_state, "embedding.model_name", - mode="file", path_modifier=components.json_path_modifier + self.components.path_entry( + frame, 0, 1, ui_state, "embedding.model_name", + mode="file", path_modifier=path_util.json_path_modifier ) # token count - components.label(frame, 1, 0, "Token count", + self.components.label(frame, 1, 0, "Token count", tooltip="The token count used when creating a new embedding. Leave empty to auto detect from the initial embedding text.") - components.entry(frame, 1, 1, self.ui_state, "embedding.token_count") + self.components.entry(frame, 1, 1, ui_state, "embedding.token_count") # initial embedding text - components.label(frame, 2, 0, "Initial embedding text", + self.components.label(frame, 2, 0, "Initial embedding text", tooltip="The initial embedding text used when creating a new embedding") - components.entry(frame, 2, 1, self.ui_state, "embedding.initial_embedding_text") + self.components.entry(frame, 2, 1, ui_state, "embedding.initial_embedding_text") # embedding weight dtype - components.label(frame, 3, 0, "Embedding Weight Data Type", + self.components.label(frame, 3, 0, "Embedding Weight Data Type", tooltip="The Embedding weight data type used for training. This can reduce memory consumption, but reduces precision") - components.options_kv(frame, 3, 1, [ + self.components.options_kv(frame, 3, 1, [ ("float32", DataType.FLOAT_32), ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "embedding_weight_dtype") + ], ui_state, "embedding_weight_dtype") # placeholder - components.label(frame, 4, 0, "Placeholder", + self.components.label(frame, 4, 0, "Placeholder", tooltip="The placeholder used when using the embedding in a prompt") - components.entry(frame, 4, 1, self.ui_state, "embedding.placeholder") + self.components.entry(frame, 4, 1, ui_state, "embedding.placeholder") # output embedding - components.label(frame, 5, 0, "Output embedding", + self.components.label(frame, 5, 0, "Output embedding", tooltip="Output embeddings are calculated at the output of the text encoder, not the input. This can improve results for larger text encoders and lower VRAM usage.") - components.switch(frame, 5, 1, self.ui_state, "embedding.is_output_embedding") - - frame.pack(fill="both", expand=1) - return frame - - def create_additional_embeddings_tab(self, master): - return AdditionalEmbeddingsTab(master, self.train_config, self.ui_state) - - def create_tools_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) + self.components.switch(frame, 5, 1, ui_state, "embedding.is_output_embedding") + def build_tools_tab_content(self, frame, controller, ui_state): # dataset - components.label(frame, 0, 0, "Dataset Tools", + self.components.label(frame, 0, 0, "Dataset Tools", tooltip="Open the captioning tool") - components.button(frame, 0, 1, "Open", self.open_dataset_tool) + self.components.button(frame, 0, 1, "Open", self.open_dataset_tool) # video tools - components.label(frame, 1, 0, "Video Tools", + self.components.label(frame, 1, 0, "Video Tools", tooltip="Open the video tools") - components.button(frame, 1, 1, "Open", self.open_video_tool) + self.components.button(frame, 1, 1, "Open", self.open_video_tool) # convert model - components.label(frame, 2, 0, "Convert Model Tools", + self.components.label(frame, 2, 0, "Convert Model Tools", tooltip="Open the model conversion tool") - components.button(frame, 2, 1, "Open", self.open_convert_model_tool) + self.components.button(frame, 2, 1, "Open", self.open_convert_model_tool) # sample - components.label(frame, 3, 0, "Sampling Tool", + self.components.label(frame, 3, 0, "Sampling Tool", tooltip="Open the model sampling tool") - components.button(frame, 3, 1, "Open", self.open_sampling_tool) + self.components.button(frame, 3, 1, "Open", self.open_sampling_tool) - components.label(frame, 4, 0, "Profiling Tool", + self.components.label(frame, 4, 0, "Profiling Tool", tooltip="Open the profiling tools.") - components.button(frame, 4, 1, "Open", self.open_profiling_tool) - - frame.pack(fill="both", expand=1) - return frame - - def change_model_type(self, model_type: ModelType): - if self.model_tab: - self.model_tab.refresh_ui() - - if self.training_tab: - self.training_tab.refresh_ui() - - if self.lora_tab: - self.lora_tab.refresh_ui() - - def change_training_method(self, training_method: TrainingMethod): - if not self.tabview: - return - - if self.model_tab: - self.model_tab.refresh_ui() - - if training_method != TrainingMethod.LORA and "LoRA" in self.tabview._tab_dict: - self.tabview.delete("LoRA") - self.lora_tab = None - if training_method != TrainingMethod.EMBEDDING and "embedding" in self.tabview._tab_dict: - self.tabview.delete("embedding") - - if training_method == TrainingMethod.LORA and "LoRA" not in self.tabview._tab_dict: - self.lora_tab = LoraTab(self.tabview.add("LoRA"), self.train_config, self.ui_state) - if training_method == TrainingMethod.EMBEDDING and "embedding" not in self.tabview._tab_dict: - self.embedding_tab(self.tabview.add("embedding")) - - def load_preset(self): - if not self.tabview: - return - - if self.additional_embeddings_tab: - self.additional_embeddings_tab.refresh_ui() - - def open_tensorboard(self): - webbrowser.open("http://localhost:" + str(self.train_config.tensorboard_port), new=0, autoraise=False) - - def _calculate_eta_string(self, train_progress: TrainProgress, max_step: int, max_epoch: int) -> str | None: - spent_total = time.monotonic() - self.start_time - steps_done = train_progress.epoch * max_step + train_progress.epoch_step - remaining_steps = (max_epoch - train_progress.epoch - 1) * max_step + (max_step - train_progress.epoch_step) - total_eta = spent_total / steps_done * remaining_steps - - if train_progress.global_step <= 30: - return "Estimating ..." - - td = datetime.timedelta(seconds=total_eta) - days = td.days - hours, remainder = divmod(td.seconds, 3600) - minutes, seconds = divmod(remainder, 60) - if days > 0: - return f"{days}d {hours}h" - elif hours > 0: - return f"{hours}h {minutes}m" - elif minutes > 0: - return f"{minutes}m {seconds}s" - else: - return f"{seconds}s" - - def set_eta_label(self, train_progress: TrainProgress, max_step: int, max_epoch: int): - eta_str = self._calculate_eta_string(train_progress, max_step, max_epoch) - if eta_str is not None: - self.eta_label.configure(text=f"ETA: {eta_str}") - else: - self.eta_label.configure(text="") - - def delete_eta_label(self): - self.eta_label.configure(text="") - - def on_update_train_progress(self, train_progress: TrainProgress, max_step: int, max_epoch: int): - self.set_step_progress(train_progress.epoch_step, max_step) - self.set_epoch_progress(train_progress.epoch, max_epoch) - self.set_eta_label(train_progress, max_step, max_epoch) - - def on_update_status(self, status: str): - self.status_label.configure(text=status) - - def open_dataset_tool(self): - window = CaptionUI(self, None, False) - self.wait_window(window) - - def open_video_tool(self): - window = VideoToolUI(self) - self.wait_window(window) - - def open_convert_model_tool(self): - window = ConvertModelUI(self) - self.wait_window(window) - - def open_sampling_tool(self): - if not self.training_callbacks and not self.training_commands: - window = SampleWindow( - self, - use_external_model=False, - train_config=self.train_config, - ) - self.wait_window(window) - torch_gc() - - def open_profiling_tool(self): - self.profiling_window.deiconify() - - def generate_debug_package(self): - zip_path = filedialog.askdirectory( - initialdir=".", - title="Select Directory to Save Debug Package" - ) - - if not zip_path: - return - - zip_path = Path(zip_path) / "OneTrainer_debug_report.zip" - - self.on_update_status("Generating debug package...") - - try: - config_json_string = json.dumps(self.train_config.to_pack_dict(secrets=False)) - scripts.generate_debug_report.create_debug_package(str(zip_path), config_json_string) - self.on_update_status(f"Debug package saved to {zip_path.name}") - except Exception as e: - traceback.print_exc() - self.on_update_status(f"Error generating debug package: {e}") - - - def open_manual_sample_window (self): - training_callbacks = self.training_callbacks - training_commands = self.training_commands - - if training_callbacks and training_commands: - window = SampleWindow( - self, - train_config=self.train_config, - use_external_model=True, - callbacks=training_callbacks, - commands=training_commands, - ) - self.wait_window(window) - training_callbacks.set_on_sample_custom() - - def __training_thread_function(self): - error_caught = False - - self.training_callbacks = TrainCallbacks( - on_update_train_progress=self.on_update_train_progress, - on_update_status=self.on_update_status, - ) - - trainer = create.create_trainer(self.train_config, self.training_callbacks, self.training_commands, reattach=self.cloud_tab.reattach) - try: - trainer.start() - if self.train_config.cloud.enabled: - self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) - - self.start_time = time.monotonic() - trainer.train() - except Exception: - if self.train_config.cloud.enabled: - self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) - error_caught = True - traceback.print_exc() - - trainer.end() - - # clear gpu memory - del trainer - - 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() - - # 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) - - def start_training(self): - if self.training_thread is None: - self.save_default() - - # --- pre-training validation gate --- - errors = flush_and_validate_all() - - if errors: - bullet_list = "\n".join(f"• {e}" for e in errors) - messagebox.showerror( - "Cannot Start Training", - f"Please fix the following errors before training:\n\n{bullet_list}", - ) - return - - 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() - - self.training_thread = threading.Thread(target=self.__training_thread_function) - self.training_thread.start() - else: - self._set_training_button_stopping() - self.on_update_status("Stopping ...") - self.training_commands.stop() - - 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() - - def export_training(self): - file_path = filedialog.asksaveasfilename(filetypes=[ - ("All Files", "*.*"), - ("json", "*.json"), - ], initialdir=".", initialfile="config.json") - - if file_path: - with open(file_path, "w") as f: - json.dump(self.train_config.to_pack_dict(secrets=False), f, indent=4) - - def sample_now(self): - train_commands = self.training_commands - if train_commands: - train_commands.sample_default() - - def backup_now(self): - train_commands = self.training_commands - if train_commands: - train_commands.backup() - - def save_now(self): - train_commands = self.training_commands - 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") - - 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", - ] - - if self.train_config.tensorboard_expose: - tensorboard_args.append("--bind_all") - - 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 _set_training_button_style(self, mode: str): - if not self.training_button: - return - style = self._TRAIN_BUTTON_STYLES.get(mode) - if not style: - return - self.training_button.configure(**style) - - def _set_training_button_idle(self): - self._set_training_button_style("idle") - - def _set_training_button_running(self): - self._set_training_button_style("running") - - def _set_training_button_stopping(self): - self._set_training_button_style("stopping") + self.components.button(frame, 4, 1, "Open", self.open_profiling_tool) diff --git a/modules/ui/BaseTrainingTabView.py b/modules/ui/BaseTrainingTabView.py index bcca11ae9..22a6af34e 100644 --- a/modules/ui/BaseTrainingTabView.py +++ b/modules/ui/BaseTrainingTabView.py @@ -1,9 +1,5 @@ -from modules.ui.OffloadingWindow import OffloadingWindow -from modules.ui.OptimizerParamsWindow import OptimizerParamsWindow -from modules.ui.SchedulerParamsWindow import SchedulerParamsWindow -from modules.ui.TimestepDistributionWindow import TimestepDistributionWindow -from modules.util import create -from modules.util.config.TrainConfig import TrainConfig +from abc import ABC + from modules.util.enum.DataType import DataType from modules.util.enum.EMAMode import EMAMode from modules.util.enum.GradientCheckpointingMethod import GradientCheckpointingMethod @@ -13,844 +9,750 @@ from modules.util.enum.LossWeight import LossWeight from modules.util.enum.Optimizer import Optimizer from modules.util.enum.TimestepDistribution import TimestepDistribution -from modules.util.optimizer_util import change_optimizer -from modules.util.ui import components -from modules.util.ui.UIState import UIState from modules.util.ui.validation_helpers import check_range, validate_resolution -import customtkinter as ctk - - -class TrainingTab: - - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__() - - self.master = master - self.train_config = train_config - self.ui_state = ui_state - - master.grid_rowconfigure(0, weight=1) - master.grid_columnconfigure(0, weight=1) - - self.scroll_frame = None - - self.refresh_ui() - - def refresh_ui(self): - if self.scroll_frame: - self.scroll_frame.destroy() - - self.scroll_frame = ctk.CTkScrollableFrame(self.master, fg_color="transparent") - self.scroll_frame.grid(row=0, column=0, sticky="nsew") - - self.scroll_frame.grid_columnconfigure(0, weight=1) - self.scroll_frame.grid_columnconfigure(1, weight=1) - self.scroll_frame.grid_columnconfigure(2, weight=1) - - column_0 = ctk.CTkFrame(master=self.scroll_frame, corner_radius=0, fg_color="transparent") - column_0.grid(row=0, column=0, sticky="nsew") - column_0.grid_columnconfigure(0, weight=1) - - column_1 = ctk.CTkFrame(master=self.scroll_frame, corner_radius=0, fg_color="transparent") - column_1.grid(row=0, column=1, sticky="nsew") - column_1.grid_columnconfigure(0, weight=1) - - column_2 = ctk.CTkFrame(master=self.scroll_frame, corner_radius=0, fg_color="transparent") - column_2.grid(row=0, column=2, sticky="nsew") - column_2.grid_columnconfigure(0, weight=1) - - if self.train_config.model_type.is_stable_diffusion(): - self.__setup_stable_diffusion_ui(column_0, column_1, column_2) - if self.train_config.model_type.is_stable_diffusion_3(): - self.__setup_stable_diffusion_3_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_stable_diffusion_xl(): - self.__setup_stable_diffusion_xl_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_wuerstchen(): - self.__setup_wuerstchen_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_pixart(): - self.__setup_pixart_alpha_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_flux_1(): - self.__setup_flux_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_flux_2(): - self.__setup_flux_2_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_chroma(): - self.__setup_chroma_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_qwen(): - self.__setup_qwen_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_sana(): - self.__setup_sana_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_hunyuan_video(): - self.__setup_hunyuan_video_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_hi_dream(): - self.__setup_hi_dream_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_z_image(): - self.__setup_z_image_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_ernie(): - self.__setup_ernie_ui(column_0, column_1, column_2) - - - def __setup_stable_diffusion_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1) - self.__create_embedding_frame(column_0, 2) - - self.__create_base2_frame(column_1, 0, supports_circular_padding=True) - self.__create_unet_frame(column_1, 1) - self.__create_noise_frame(column_1, 2, supports_generalized_offset_noise=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_stable_diffusion_3_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 3, i=3, supports_include=True) - self.__create_embedding_frame(column_0, 4) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_stable_diffusion_xl_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_n_frame(column_0, 1, i=1) - self.__create_text_encoder_n_frame(column_0, 2, i=2) - self.__create_embedding_frame(column_0, 3) - - self.__create_base2_frame(column_1, 0, supports_circular_padding=True) - self.__create_unet_frame(column_1, 1) - self.__create_noise_frame(column_1, 2, supports_generalized_offset_noise=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_wuerstchen_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1) - self.__create_embedding_frame(column_0, 2) - - self.__create_base2_frame(column_1, 0, supports_circular_padding=True) - self.__create_prior_frame(column_1, 1) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 0) - self.__create_loss_frame(column_2, 1) - self.__create_layer_frame(column_2, 2) - - def __setup_pixart_alpha_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1) - self.__create_embedding_frame(column_0, 2) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2, supports_vb_loss=True) - self.__create_layer_frame(column_2, 3) - - def __setup_flux_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True, supports_sequence_length=True) - self.__create_embedding_frame(column_0, 4) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=True) - self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_flux_2_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False, supports_sequence_length=True) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=True, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_chroma_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1) - self.__create_embedding_frame(column_0, 4) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_qwen_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_z_image_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_ernie_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_sana_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1) - self.__create_embedding_frame(column_0, 2) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_hunyuan_video_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True) - self.__create_embedding_frame(column_0, 4) - - self.__create_base2_frame(column_1, 0, video_training_enabled=True) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=True) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_hi_dream_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 3, i=3, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 4, i=4, supports_include=True, supports_layer_skip=False) - self.__create_embedding_frame(column_0, 5) - - self.__create_base2_frame(column_1, 0, video_training_enabled=True) - self.__create_transformer_frame(column_1, 1) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __create_base_frame(self, master, row): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) + +class BaseTrainingTabView(ABC): + def __init__(self, components): + self.components = components + + def build(self, column_0, column_1, column_2, controller, ui_state, callbacks: dict): + model_type = controller.config.model_type + if model_type.is_stable_diffusion(): + self.__setup_stable_diffusion_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + if model_type.is_stable_diffusion_3(): + self.__setup_stable_diffusion_3_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_stable_diffusion_xl(): + self.__setup_stable_diffusion_xl_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_wuerstchen(): + self.__setup_wuerstchen_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_pixart(): + self.__setup_pixart_alpha_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_flux_1(): + self.__setup_flux_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_flux_2(): + self.__setup_flux_2_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_chroma(): + self.__setup_chroma_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_qwen(): + self.__setup_qwen_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_sana(): + self.__setup_sana_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_hunyuan_video(): + self.__setup_hunyuan_video_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_hi_dream(): + self.__setup_hi_dream_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_z_image(): + self.__setup_z_image_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_ernie(): + self.__setup_ernie_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + + def __setup_stable_diffusion_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_frame(column_0, 1, ui_state) + self.__create_embedding_frame(column_0, 2, ui_state) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks, supports_circular_padding=True) + self.__create_unet_frame(column_1, 1, ui_state) + self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_generalized_offset_noise=True) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_stable_diffusion_3_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_n_frame(column_0, 1, ui_state, i=1, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 2, ui_state, i=2, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 3, ui_state, i=3, supports_include=True) + self.__create_embedding_frame(column_0, 4, ui_state) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_transformer_frame(column_1, 1, ui_state) + self.__create_noise_frame(column_1, 2, ui_state, callbacks) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_stable_diffusion_xl_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_n_frame(column_0, 1, ui_state, i=1) + self.__create_text_encoder_n_frame(column_0, 2, ui_state, i=2) + self.__create_embedding_frame(column_0, 3, ui_state) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks, supports_circular_padding=True) + self.__create_unet_frame(column_1, 1, ui_state) + self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_generalized_offset_noise=True) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_wuerstchen_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_frame(column_0, 1, ui_state) + self.__create_embedding_frame(column_0, 2, ui_state) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks, supports_circular_padding=True) + self.__create_prior_frame(column_1, 1, ui_state) + self.__create_noise_frame(column_1, 2, ui_state, callbacks) + + self.__create_masked_frame(column_2, 0, ui_state) + self.__create_loss_frame(column_2, 1, controller, ui_state) + self.__create_layer_frame(column_2, 2, controller, ui_state) + + def __setup_pixart_alpha_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_frame(column_0, 1, ui_state) + self.__create_embedding_frame(column_0, 2, ui_state) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_transformer_frame(column_1, 1, ui_state) + self.__create_noise_frame(column_1, 2, ui_state, callbacks) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state, supports_vb_loss=True) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_flux_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_n_frame(column_0, 1, ui_state, i=1, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 2, ui_state, i=2, supports_include=True, supports_sequence_length=True) + self.__create_embedding_frame(column_0, 4, ui_state) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=True) + self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_flux_2_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_frame(column_0, 1, ui_state, supports_clip_skip=False, supports_training=False, supports_sequence_length=True) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=True, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_chroma_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_frame(column_0, 1, ui_state) + self.__create_embedding_frame(column_0, 4, ui_state) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=False, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, ui_state, callbacks) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_qwen_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_frame(column_0, 1, ui_state, supports_clip_skip=False) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=False, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_z_image_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_frame(column_0, 1, ui_state, supports_clip_skip=False, supports_training=False) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=False, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_ernie_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_frame(column_0, 1, ui_state, supports_clip_skip=False, supports_training=False) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=False, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_sana_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_frame(column_0, 1, ui_state) + self.__create_embedding_frame(column_0, 2, ui_state) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_transformer_frame(column_1, 1, ui_state) + self.__create_noise_frame(column_1, 2, ui_state, callbacks) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_hunyuan_video_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_n_frame(column_0, 1, ui_state, i=1, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 2, ui_state, i=2, supports_include=True) + self.__create_embedding_frame(column_0, 4, ui_state) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks, video_training_enabled=True) + self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=True) + self.__create_noise_frame(column_1, 2, ui_state, callbacks) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_hi_dream_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_n_frame(column_0, 1, ui_state, i=1, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 2, ui_state, i=2, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 3, ui_state, i=3, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 4, ui_state, i=4, supports_include=True, supports_layer_skip=False) + self.__create_embedding_frame(column_0, 5, ui_state) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks, video_training_enabled=True) + self.__create_transformer_frame(column_1, 1, ui_state) + self.__create_noise_frame(column_1, 2, ui_state, callbacks) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __create_base_frame(self, master, row, controller, ui_state, callbacks): + frame = self.components.section_frame(master, row) # optimizer - components.label(frame, 0, 0, "Optimizer", - tooltip="The type of optimizer") - components.options_adv(frame, 0, 1, [str(x) for x in list(Optimizer)], self.ui_state, "optimizer.optimizer", - command=self.__restore_optimizer_config, adv_command=self.__open_optimizer_params_window) + self.components.label(frame, 0, 0, "Optimizer", + tooltip="The type of optimizer") + self.components.options_adv(frame, 0, 1, [str(x) for x in list(Optimizer)], ui_state, "optimizer.optimizer", + command=callbacks.get('restore_optimizer'), + adv_command=callbacks.get('open_optimizer_params')) # learning rate scheduler # Wackiness will ensue when reloading configs if we don't check and clear this first. if hasattr(self, "lr_scheduler_comp"): delattr(self, "lr_scheduler_comp") delattr(self, "lr_scheduler_adv_comp") - components.label(frame, 1, 0, "Learning Rate Scheduler", - tooltip="Learning rate scheduler that automatically changes the learning rate during training") - _, d = components.options_adv(frame, 1, 1, [str(x) for x in list(LearningRateScheduler)], self.ui_state, - "learning_rate_scheduler", command=self.__restore_scheduler_config, - adv_command=self.__open_scheduler_params_window) + self.components.label(frame, 1, 0, "Learning Rate Scheduler", + tooltip="Learning rate scheduler that automatically changes the learning rate during training") + _, d = self.components.options_adv(frame, 1, 1, [str(x) for x in list(LearningRateScheduler)], ui_state, + "learning_rate_scheduler", + command=callbacks.get('restore_scheduler'), + adv_command=callbacks.get('open_scheduler_params')) self.lr_scheduler_comp = d['component'] self.lr_scheduler_adv_comp = d['button_component'] # Initial call requires the presence of self.lr_scheduler_adv_comp. - self.__restore_scheduler_config(self.ui_state.get_var("learning_rate_scheduler").get()) + restore_scheduler = callbacks.get('restore_scheduler') + if restore_scheduler: + restore_scheduler(ui_state.get_var("learning_rate_scheduler").get()) # learning rate - components.label(frame, 2, 0, "Learning Rate", - tooltip="The base learning rate") - components.entry(frame, 2, 1, self.ui_state, "learning_rate", required=True) + self.components.label(frame, 2, 0, "Learning Rate", + tooltip="The base learning rate") + self.components.entry(frame, 2, 1, ui_state, "learning_rate", required=True) # learning rate warmup steps - components.label(frame, 3, 0, "Learning Rate Warmup Steps", - tooltip="The number of steps it takes to gradually increase the learning rate from 0 to the specified learning rate. Values >1 are interpeted as a fixed number of steps, values <=1 are intepreted as a percentage of the total training steps (ex. 0.2 = 20% of the total step count)") - components.entry(frame, 3, 1, self.ui_state, "learning_rate_warmup_steps") + self.components.label(frame, 3, 0, "Learning Rate Warmup Steps", + tooltip="The number of steps it takes to gradually increase the learning rate from 0 to the specified learning rate. Values >1 are interpeted as a fixed number of steps, values <=1 are intepreted as a percentage of the total training steps (ex. 0.2 = 20% of the total step count)") + self.components.entry(frame, 3, 1, ui_state, "learning_rate_warmup_steps") # learning rate min factor - components.label(frame, 4, 0, "Learning Rate Min Factor", - tooltip="Unit = float. Method = percentage. For a factor of 0.1, the final LR will be 10% of the initial LR. If the initial LR is 1e-4, the final LR will be 1e-5.") - components.entry(frame, 4, 1, self.ui_state, "learning_rate_min_factor", - extra_validate=check_range(lower=0, upper=0.99, message="Learning rate min factor must be between 0 and 0.99")) + self.components.label(frame, 4, 0, "Learning Rate Min Factor", + tooltip="Unit = float. Method = percentage. For a factor of 0.1, the final LR will be 10% of the initial LR. If the initial LR is 1e-4, the final LR will be 1e-5.") + self.components.entry(frame, 4, 1, ui_state, "learning_rate_min_factor", + extra_validate=check_range(lower=0, upper=0.99, message="Learning rate min factor must be between 0 and 0.99")) # learning rate cycles - components.label(frame, 5, 0, "Learning Rate Cycles", - tooltip="The number of learning rate cycles. This is only applicable if the learning rate scheduler supports cycles") - components.entry(frame, 5, 1, self.ui_state, "learning_rate_cycles") + self.components.label(frame, 5, 0, "Learning Rate Cycles", + tooltip="The number of learning rate cycles. This is only applicable if the learning rate scheduler supports cycles") + self.components.entry(frame, 5, 1, ui_state, "learning_rate_cycles") # epochs - components.label(frame, 6, 0, "Epochs", - tooltip="The number of epochs for a full training run") - components.entry(frame, 6, 1, self.ui_state, "epochs", required=True) + self.components.label(frame, 6, 0, "Epochs", + tooltip="The number of epochs for a full training run") + self.components.entry(frame, 6, 1, ui_state, "epochs", required=True) # batch size - components.label(frame, 7, 0, "Local Batch Size", - tooltip="The batch size of one training step. If you use multiple GPUs, this is the batch size of each GPU (local batch size).") - components.entry(frame, 7, 1, self.ui_state, "batch_size", required=True) + self.components.label(frame, 7, 0, "Local Batch Size", + tooltip="The batch size of one training step. If you use multiple GPUs, this is the batch size of each GPU (local batch size).") + self.components.entry(frame, 7, 1, ui_state, "batch_size", required=True) # accumulation steps - components.label(frame, 8, 0, "Accumulation Steps", - tooltip="Number of accumulation steps. Increase this number to trade batch size for training speed") - components.entry(frame, 8, 1, self.ui_state, "gradient_accumulation_steps", required=True) + self.components.label(frame, 8, 0, "Accumulation Steps", + tooltip="Number of accumulation steps. Increase this number to trade batch size for training speed") + self.components.entry(frame, 8, 1, ui_state, "gradient_accumulation_steps", required=True) # Learning Rate Scaler - components.label(frame, 9, 0, "Learning Rate Scaler", - tooltip="Selects the type of learning rate scaling to use during training. Functionally equated as: LR * SQRT(selection)") - components.options(frame, 9, 1, [str(x) for x in list(LearningRateScaler)], self.ui_state, - "learning_rate_scaler") + self.components.label(frame, 9, 0, "Learning Rate Scaler", + tooltip="Selects the type of learning rate scaling to use during training. Functionally equated as: LR * SQRT(selection)") + self.components.options(frame, 9, 1, [str(x) for x in list(LearningRateScaler)], ui_state, + "learning_rate_scaler") # clip grad norm - components.label(frame, 10, 0, "Clip Grad Norm", - tooltip="Clips the gradient norm. Leave empty to disable gradient clipping.") - components.entry(frame, 10, 1, self.ui_state, "clip_grad_norm") - - def __create_base2_frame(self, master, row, video_training_enabled: bool=False, supports_circular_padding: bool=False): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) + self.components.label(frame, 10, 0, "Clip Grad Norm", + tooltip="Clips the gradient norm. Leave empty to disable gradient clipping.") + self.components.entry(frame, 10, 1, ui_state, "clip_grad_norm") + + def __create_base2_frame(self, master, row, ui_state, callbacks, video_training_enabled: bool = False, + supports_circular_padding: bool = False): + frame = self.components.section_frame(master, row) row = 0 # ema - components.label(frame, row, 0, "EMA", - tooltip="EMA averages the training progress over many steps, better preserving different concepts in big datasets") - components.options(frame, row, 1, [str(x) for x in list(EMAMode)], self.ui_state, "ema") + self.components.label(frame, row, 0, "EMA", + tooltip="EMA averages the training progress over many steps, better preserving different concepts in big datasets") + self.components.options(frame, row, 1, [str(x) for x in list(EMAMode)], ui_state, "ema") row += 1 # ema decay - components.label(frame, row, 0, "EMA Decay", - tooltip="Decay parameter of the EMA model. Higher numbers will average more steps. For datasets of hundreds or thousands of images, set this to 0.9999. For smaller datasets, set it to 0.999 or even 0.998") - components.entry(frame, row, 1, self.ui_state, "ema_decay", - extra_validate=check_range(lower=0.5, upper=1, - message="EMA decay must be between 0.5 and 1")) + self.components.label(frame, row, 0, "EMA Decay", + tooltip="Decay parameter of the EMA model. Higher numbers will average more steps. For datasets of hundreds or thousands of images, set this to 0.9999. For smaller datasets, set it to 0.999 or even 0.998") + self.components.entry(frame, row, 1, ui_state, "ema_decay", + extra_validate=check_range(lower=0.5, upper=1, + message="EMA decay must be between 0.5 and 1")) row += 1 # ema update step interval - components.label(frame, row, 0, "EMA Update Step Interval", - tooltip="Number of steps between EMA update steps") - components.entry(frame, row, 1, self.ui_state, "ema_update_step_interval") + self.components.label(frame, row, 0, "EMA Update Step Interval", + tooltip="Number of steps between EMA update steps") + self.components.entry(frame, row, 1, ui_state, "ema_update_step_interval") row += 1 # gradient checkpointing - components.label(frame, row, 0, "Gradient checkpointing", - tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") - components.options_adv(frame, row, 1, [str(x) for x in list(GradientCheckpointingMethod)], self.ui_state, - "gradient_checkpointing", adv_command=self.__open_offloading_window) + self.components.label(frame, row, 0, "Gradient checkpointing", + tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") + self.components.options_adv(frame, row, 1, [str(x) for x in list(GradientCheckpointingMethod)], ui_state, + "gradient_checkpointing", + adv_command=callbacks.get('open_offloading')) row += 1 # gradient checkpointing layer offloading - components.label(frame, row, 0, "Layer offload fraction", - tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") - components.entry(frame, row, 1, self.ui_state, "layer_offload_fraction") + self.components.label(frame, row, 0, "Layer offload fraction", + tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") + self.components.entry(frame, row, 1, ui_state, "layer_offload_fraction") row += 1 # train dtype - components.label(frame, row, 0, "Train Data Type", - tooltip="The mixed precision data type used for training. This can increase training speed, but reduces precision") - components.options_kv(frame, row, 1, [ + self.components.label(frame, row, 0, "Train Data Type", + tooltip="The mixed precision data type used for training. This can increase training speed, but reduces precision") + self.components.options_kv(frame, row, 1, [ ("float32", DataType.FLOAT_32), ("float16", DataType.FLOAT_16), ("bfloat16", DataType.BFLOAT_16), ("tfloat32", DataType.TFLOAT_32), - ], self.ui_state, "train_dtype") + ], ui_state, "train_dtype") row += 1 # fallback train dtype - components.label(frame, row, 0, "Fallback Train Data Type", - tooltip="The mixed precision data type used for training stages that don't support float16 data types. This can increase training speed, but reduces precision") - components.options_kv(frame, row, 1, [ + self.components.label(frame, row, 0, "Fallback Train Data Type", + tooltip="The mixed precision data type used for training stages that don't support float16 data types. This can increase training speed, but reduces precision") + self.components.options_kv(frame, row, 1, [ ("float32", DataType.FLOAT_32), ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "fallback_train_dtype") + ], ui_state, "fallback_train_dtype") row += 1 # autocast cache - components.label(frame, row, 0, "Autocast Cache", - tooltip="Enables the autocast cache. Disabling this reduces memory usage, but increases training time") - components.switch(frame, row, 1, self.ui_state, "enable_autocast_cache") + self.components.label(frame, row, 0, "Autocast Cache", + tooltip="Enables the autocast cache. Disabling this reduces memory usage, but increases training time") + self.components.switch(frame, row, 1, ui_state, "enable_autocast_cache") row += 1 # resolution - components.label(frame, row, 0, "Resolution", - tooltip="The resolution used for training. Optionally specify multiple resolutions separated by a comma, or a single exact resolution in the format x") - components.entry(frame, row, 1, self.ui_state, "resolution", required=True, - extra_validate=validate_resolution()) + self.components.label(frame, row, 0, "Resolution", + tooltip="The resolution used for training. Optionally specify multiple resolutions separated by a comma, or a single exact resolution in the format x") + self.components.entry(frame, row, 1, ui_state, "resolution", required=True, + extra_validate=validate_resolution()) row += 1 # frames if video_training_enabled: - components.label(frame, row, 0, "Frames", - tooltip="The number of frames used for training.") - components.entry(frame, row, 1, self.ui_state, "frames", required=True) + self.components.label(frame, row, 0, "Frames", + tooltip="The number of frames used for training.") + self.components.entry(frame, row, 1, ui_state, "frames", required=True) row += 1 # force circular padding if supports_circular_padding: - components.label(frame, row, 0, "Force Circular Padding", - tooltip="Enables circular padding for all conv layers to better train seamless images") - components.switch(frame, row, 1, self.ui_state, "force_circular_padding") + self.components.label(frame, row, 0, "Force Circular Padding", + tooltip="Enables circular padding for all conv layers to better train seamless images") + self.components.switch(frame, row, 1, ui_state, "force_circular_padding") - def __create_text_encoder_frame(self, master, row, supports_clip_skip=True, supports_training=True, supports_sequence_length=False): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) + def __create_text_encoder_frame(self, master, row, ui_state, supports_clip_skip=True, supports_training=True, + supports_sequence_length=False): + frame = self.components.section_frame(master, row) if supports_training: - components.label(frame, 0, 0, "Train Text Encoder", - tooltip="Enables training the text encoder model") - components.switch(frame, 0, 1, self.ui_state, "text_encoder.train") + self.components.label(frame, 0, 0, "Train Text Encoder", + tooltip="Enables training the text encoder model") + self.components.switch(frame, 0, 1, ui_state, "text_encoder.train") # dropout - components.label(frame, 1, 0, "Caption Dropout Probability", - tooltip="The Probability for dropping the text encoder conditioning") - components.entry(frame, 1, 1, self.ui_state, "text_encoder.dropout_probability") + self.components.label(frame, 1, 0, "Caption Dropout Probability", + tooltip="The Probability for dropping the text encoder conditioning") + self.components.entry(frame, 1, 1, ui_state, "text_encoder.dropout_probability") if supports_training: # train text encoder epochs - components.label(frame, 2, 0, "Stop Training After", - tooltip="When to stop training the text encoder") - components.time_entry(frame, 2, 1, self.ui_state, "text_encoder.stop_training_after", - "text_encoder.stop_training_after_unit", supports_time_units=False) + self.components.label(frame, 2, 0, "Stop Training After", + tooltip="When to stop training the text encoder") + self.components.time_entry(frame, 2, 1, ui_state, "text_encoder.stop_training_after", + "text_encoder.stop_training_after_unit", supports_time_units=False) # text encoder learning rate - components.label(frame, 3, 0, "Text Encoder Learning Rate", - tooltip="The learning rate of the text encoder. Overrides the base learning rate") - components.entry(frame, 3, 1, self.ui_state, "text_encoder.learning_rate") + self.components.label(frame, 3, 0, "Text Encoder Learning Rate", + tooltip="The learning rate of the text encoder. Overrides the base learning rate") + self.components.entry(frame, 3, 1, ui_state, "text_encoder.learning_rate") if supports_clip_skip: # text encoder layer skip (clip skip) - components.label(frame, 4, 0, "Clip Skip", - tooltip="The number of additional clip layers to skip. 0 = the model default") - components.entry(frame, 4, 1, self.ui_state, "text_encoder_layer_skip") + self.components.label(frame, 4, 0, "Clip Skip", + tooltip="The number of additional clip layers to skip. 0 = the model default") + self.components.entry(frame, 4, 1, ui_state, "text_encoder_layer_skip") if supports_sequence_length: # text encoder sequence length - components.label(frame, row, 0, "Text Encoder Sequence Length", - tooltip="Number of tokens for captions") - components.entry(frame, row, 1, self.ui_state, "text_encoder_sequence_length") + self.components.label(frame, row, 0, "Text Encoder Sequence Length", + tooltip="Number of tokens for captions") + self.components.entry(frame, row, 1, ui_state, "text_encoder_sequence_length") row += 1 def __create_text_encoder_n_frame( self, master, row: int, + ui_state, i: int, supports_include: bool = False, supports_layer_skip: bool = True, supports_sequence_length: bool = False, ): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) + frame = self.components.section_frame(master, row) row = 0 suffix = f"_{i}" if i > 1 else "" if supports_include: # include text encoder - components.label(frame, row, 0, f"Include Text Encoder {i}", - tooltip=f"Includes text encoder {i} in the training run") - components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.include") + self.components.label(frame, row, 0, f"Include Text Encoder {i}", + tooltip=f"Includes text encoder {i} in the training run") + self.components.switch(frame, row, 1, ui_state, f"text_encoder{suffix}.include") row += 1 # train text encoder - components.label(frame, row, 0, f"Train Text Encoder {i}", - tooltip=f"Enables training the text encoder {i} model") - components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.train") + self.components.label(frame, row, 0, f"Train Text Encoder {i}", + tooltip=f"Enables training the text encoder {i} model") + self.components.switch(frame, row, 1, ui_state, f"text_encoder{suffix}.train") row += 1 # train text encoder embedding - components.label(frame, row, 0, f"Train Text Encoder {i} Embedding", - tooltip=f"Enables training embeddings for the text encoder {i} model") - components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.train_embedding") + self.components.label(frame, row, 0, f"Train Text Encoder {i} Embedding", + tooltip=f"Enables training embeddings for the text encoder {i} model") + self.components.switch(frame, row, 1, ui_state, f"text_encoder{suffix}.train_embedding") row += 1 # dropout - components.label(frame, row, 0, "Dropout Probability", - tooltip=f"The Probability for dropping the text encoder {i} conditioning") - components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}.dropout_probability") + self.components.label(frame, row, 0, "Dropout Probability", + tooltip=f"The Probability for dropping the text encoder {i} conditioning") + self.components.entry(frame, row, 1, ui_state, f"text_encoder{suffix}.dropout_probability") row += 1 # train text encoder epochs - components.label(frame, row, 0, "Stop Training After", - tooltip=f"When to stop training the text encoder {i}") - components.time_entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}.stop_training_after", - f"text_encoder{suffix}.stop_training_after_unit", supports_time_units=False) + self.components.label(frame, row, 0, "Stop Training After", + tooltip=f"When to stop training the text encoder {i}") + self.components.time_entry(frame, row, 1, ui_state, f"text_encoder{suffix}.stop_training_after", + f"text_encoder{suffix}.stop_training_after_unit", supports_time_units=False) row += 1 # text encoder learning rate - components.label(frame, row, 0, f"Text Encoder {i} Learning Rate", - tooltip=f"The learning rate of the text encoder {i}. Overrides the base learning rate") - components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}.learning_rate") + self.components.label(frame, row, 0, f"Text Encoder {i} Learning Rate", + tooltip=f"The learning rate of the text encoder {i}. Overrides the base learning rate") + self.components.entry(frame, row, 1, ui_state, f"text_encoder{suffix}.learning_rate") row += 1 if supports_layer_skip: # text encoder layer skip (clip skip) - components.label(frame, row, 0, f"Text Encoder {i} Clip Skip", - tooltip="The number of additional clip layers to skip. 0 = the model default") - components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}_layer_skip") + self.components.label(frame, row, 0, f"Text Encoder {i} Clip Skip", + tooltip="The number of additional clip layers to skip. 0 = the model default") + self.components.entry(frame, row, 1, ui_state, f"text_encoder{suffix}_layer_skip") row += 1 if supports_sequence_length: # text encoder sequence length - components.label(frame, row, 0, f"Text Encoder {i} Sequence Length", - tooltip="Overrides the number of tokens used for captions. If empty, the model default is used, which is 512 on Flux. Comfy samples with 256 tokens though. 77 is the default only for backwards compatibility.") - components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}_sequence_length") + self.components.label(frame, row, 0, f"Text Encoder {i} Sequence Length", + tooltip="Overrides the number of tokens used for captions. If empty, the model default is used, which is 512 on Flux. Comfy samples with 256 tokens though. 77 is the default only for backwards compatibility.") + self.components.entry(frame, row, 1, ui_state, f"text_encoder{suffix}_sequence_length") row += 1 - def __create_embedding_frame(self, master, row): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + def __create_embedding_frame(self, master, row, ui_state): + frame = self.components.section_frame(master, row) # embedding learning rate - components.label(frame, 0, 0, "Embeddings Learning Rate", - tooltip="The learning rate of embeddings. Overrides the base learning rate") - components.entry(frame, 0, 1, self.ui_state, "embedding_learning_rate") + self.components.label(frame, 0, 0, "Embeddings Learning Rate", + tooltip="The learning rate of embeddings. Overrides the base learning rate") + self.components.entry(frame, 0, 1, ui_state, "embedding_learning_rate") # preserve embedding norm - components.label(frame, 1, 0, "Preserve Embedding Norm", - tooltip="Rescales each trained embedding to the median embedding norm") - components.switch(frame, 1, 1, self.ui_state, "preserve_embedding_norm") + self.components.label(frame, 1, 0, "Preserve Embedding Norm", + tooltip="Rescales each trained embedding to the median embedding norm") + self.components.switch(frame, 1, 1, ui_state, "preserve_embedding_norm") - def __create_unet_frame(self, master, row): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) + def __create_unet_frame(self, master, row, ui_state): + frame = self.components.section_frame(master, row) # train unet - components.label(frame, 0, 0, "Train UNet", - tooltip="Enables training the UNet model") - components.switch(frame, 0, 1, self.ui_state, "unet.train") + self.components.label(frame, 0, 0, "Train UNet", + tooltip="Enables training the UNet model") + self.components.switch(frame, 0, 1, ui_state, "unet.train") # train unet epochs - components.label(frame, 1, 0, "Stop Training After", - tooltip="When to stop training the UNet") - components.time_entry(frame, 1, 1, self.ui_state, "unet.stop_training_after", "unet.stop_training_after_unit", - supports_time_units=False) + self.components.label(frame, 1, 0, "Stop Training After", + tooltip="When to stop training the UNet") + self.components.time_entry(frame, 1, 1, ui_state, "unet.stop_training_after", "unet.stop_training_after_unit", + supports_time_units=False) # unet learning rate - components.label(frame, 2, 0, "UNet Learning Rate", - tooltip="The learning rate of the UNet. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "unet.learning_rate") + self.components.label(frame, 2, 0, "UNet Learning Rate", + tooltip="The learning rate of the UNet. Overrides the base learning rate") + self.components.entry(frame, 2, 1, ui_state, "unet.learning_rate") # rescale noise scheduler to zero terminal SNR - rescale_label = components.label(frame, 3, 0, "Rescale Noise Scheduler + V-pred", - tooltip="Rescales the noise scheduler to a zero terminal signal to noise ratio and switches the model to a v-prediction target") - rescale_label.configure(wraplength=130, justify="left") - components.switch(frame, 3, 1, self.ui_state, "rescale_noise_scheduler_to_zero_terminal_snr") + self.components.label(frame, 3, 0, "Rescale Noise Scheduler + V-pred", + tooltip="Rescales the noise scheduler to a zero terminal signal to noise ratio and switches the model to a v-prediction target", + wraplength=130) + self.components.switch(frame, 3, 1, ui_state, "rescale_noise_scheduler_to_zero_terminal_snr") - def __create_prior_frame(self, master, row): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) + def __create_prior_frame(self, master, row, ui_state): + frame = self.components.section_frame(master, row) # train prior - components.label(frame, 0, 0, "Train Prior", - tooltip="Enables training the Prior model") - components.switch(frame, 0, 1, self.ui_state, "prior.train") + self.components.label(frame, 0, 0, "Train Prior", + tooltip="Enables training the Prior model") + self.components.switch(frame, 0, 1, ui_state, "prior.train") # train prior epochs - components.label(frame, 1, 0, "Stop Training After", - tooltip="When to stop training the Prior") - components.time_entry(frame, 1, 1, self.ui_state, "prior.stop_training_after", "prior.stop_training_after_unit", - supports_time_units=False) + self.components.label(frame, 1, 0, "Stop Training After", + tooltip="When to stop training the Prior") + self.components.time_entry(frame, 1, 1, ui_state, "prior.stop_training_after", + "prior.stop_training_after_unit", supports_time_units=False) # prior learning rate - components.label(frame, 2, 0, "Prior Learning Rate", - tooltip="The learning rate of the Prior. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "prior.learning_rate") + self.components.label(frame, 2, 0, "Prior Learning Rate", + tooltip="The learning rate of the Prior. Overrides the base learning rate") + self.components.entry(frame, 2, 1, ui_state, "prior.learning_rate") - def __create_transformer_frame(self, master, row, supports_guidance_scale: bool = False, supports_force_attention_mask: bool = True): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) + def __create_transformer_frame(self, master, row, ui_state, supports_guidance_scale: bool = False, + supports_force_attention_mask: bool = True): + frame = self.components.section_frame(master, row) # train transformer - components.label(frame, 0, 0, "Train Transformer", - tooltip="Enables training the Transformer model") - components.switch(frame, 0, 1, self.ui_state, "transformer.train") + self.components.label(frame, 0, 0, "Train Transformer", + tooltip="Enables training the Transformer model") + self.components.switch(frame, 0, 1, ui_state, "transformer.train") # train transformer epochs - components.label(frame, 1, 0, "Stop Training After", - tooltip="When to stop training the Transformer") - components.time_entry(frame, 1, 1, self.ui_state, "transformer.stop_training_after", "transformer.stop_training_after_unit", - supports_time_units=False) + self.components.label(frame, 1, 0, "Stop Training After", + tooltip="When to stop training the Transformer") + self.components.time_entry(frame, 1, 1, ui_state, "transformer.stop_training_after", + "transformer.stop_training_after_unit", supports_time_units=False) # transformer learning rate - components.label(frame, 2, 0, "Transformer Learning Rate", - tooltip="The learning rate of the Transformer. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "transformer.learning_rate") + self.components.label(frame, 2, 0, "Transformer Learning Rate", + tooltip="The learning rate of the Transformer. Overrides the base learning rate") + self.components.entry(frame, 2, 1, ui_state, "transformer.learning_rate") if supports_force_attention_mask: # transformer learning rate - components.label(frame, 3, 0, "Force Attention Mask", - tooltip="Force enables passing of a text embedding attention mask to the transformer. This can improve training on shorter captions.") - components.switch(frame, 3, 1, self.ui_state, "transformer.attention_mask") + self.components.label(frame, 3, 0, "Force Attention Mask", + tooltip="Force enables passing of a text embedding attention mask to the transformer. This can improve training on shorter captions.") + self.components.switch(frame, 3, 1, ui_state, "transformer.attention_mask") if supports_guidance_scale: # guidance scale - components.label(frame, 4, 0, "Guidance Scale", - tooltip="The guidance scale of guidance distilled models passed to the transformer during training.") - components.entry(frame, 4, 1, self.ui_state, "transformer.guidance_scale") + self.components.label(frame, 4, 0, "Guidance Scale", + tooltip="The guidance scale of guidance distilled models passed to the transformer during training.") + self.components.entry(frame, 4, 1, ui_state, "transformer.guidance_scale") - def __create_noise_frame(self, master, row, supports_generalized_offset_noise: bool = False, supports_dynamic_timestep_shifting: bool = False): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) + def __create_noise_frame(self, master, row, ui_state, callbacks, + supports_generalized_offset_noise: bool = False, + supports_dynamic_timestep_shifting: bool = False): + frame = self.components.section_frame(master, row) # offset noise weight - components.label(frame, 0, 0, "Offset Noise Weight", - tooltip="The weight of offset noise added to each training step") - components.entry(frame, 0, 1, self.ui_state, "offset_noise_weight") + self.components.label(frame, 0, 0, "Offset Noise Weight", + tooltip="The weight of offset noise added to each training step") + self.components.entry(frame, 0, 1, ui_state, "offset_noise_weight") if supports_generalized_offset_noise: # generalized offset noise weight - generalised_offset_label = components.label(frame, 1, 0, "Generalized Offset Noise", - tooltip="Per-timestep 'brightness knob' instead of a fixed offset - steadier training, better starts, and improved very dark/bright images. Compatible with V-pred and Eps-pred. Start with 0.02 and adjust as needed.") - generalised_offset_label.configure(wraplength=130, justify="left") - components.switch(frame, 1, 1, self.ui_state, "generalized_offset_noise") + self.components.label(frame, 1, 0, "Generalized Offset Noise", + tooltip="Per-timestep 'brightness knob' instead of a fixed offset - steadier training, better starts, and improved very dark/bright images. Compatible with V-pred and Eps-pred. Start with 0.02 and adjust as needed.", + wraplength=130) + self.components.switch(frame, 1, 1, ui_state, "generalized_offset_noise") # perturbation noise weight - components.label(frame, 2, 0, "Perturbation Noise Weight", - tooltip="The weight of perturbation noise added to each training step") - components.entry(frame, 2, 1, self.ui_state, "perturbation_noise_weight") + self.components.label(frame, 2, 0, "Perturbation Noise Weight", + tooltip="The weight of perturbation noise added to each training step") + self.components.entry(frame, 2, 1, ui_state, "perturbation_noise_weight") # timestep distribution - components.label(frame, 3, 0, "Timestep Distribution", - tooltip="Selects the function to sample timesteps during training", - wide_tooltip=True) - components.options_adv(frame, 3, 1, [str(x) for x in list(TimestepDistribution)], self.ui_state, "timestep_distribution", - adv_command=self.__open_timestep_distribution_window) + self.components.label(frame, 3, 0, "Timestep Distribution", + tooltip="Selects the function to sample timesteps during training", + wide_tooltip=True) + self.components.options_adv(frame, 3, 1, [str(x) for x in list(TimestepDistribution)], ui_state, + "timestep_distribution", + adv_command=callbacks.get('open_timestep_distribution')) # min noising strength - components.label(frame, 4, 0, "Min Noising Strength", - tooltip="Specifies the minimum noising strength used during training. This can help to improve composition, but prevents finer details from being trained") - components.entry(frame, 4, 1, self.ui_state, "min_noising_strength", required=True) + self.components.label(frame, 4, 0, "Min Noising Strength", + tooltip="Specifies the minimum noising strength used during training. This can help to improve composition, but prevents finer details from being trained") + self.components.entry(frame, 4, 1, ui_state, "min_noising_strength", required=True) # max noising strength - components.label(frame, 5, 0, "Max Noising Strength", - tooltip="Specifies the maximum noising strength used during training. This can be useful to reduce overfitting, but also reduces the impact of training samples on the overall image composition") - components.entry(frame, 5, 1, self.ui_state, "max_noising_strength", required=True) + self.components.label(frame, 5, 0, "Max Noising Strength", + tooltip="Specifies the maximum noising strength used during training. This can be useful to reduce overfitting, but also reduces the impact of training samples on the overall image composition") + self.components.entry(frame, 5, 1, ui_state, "max_noising_strength", required=True) # noising weight - components.label(frame, 6, 0, "Noising Weight", - tooltip="Controls the weight parameter of the timestep distribution function. Use the preview to see more details.") - components.entry(frame, 6, 1, self.ui_state, "noising_weight", required=True) + self.components.label(frame, 6, 0, "Noising Weight", + tooltip="Controls the weight parameter of the timestep distribution function. Use the preview to see more details.") + self.components.entry(frame, 6, 1, ui_state, "noising_weight", required=True) # noising bias - components.label(frame, 7, 0, "Noising Bias", - tooltip="Controls the bias parameter of the timestep distribution function. Use the preview to see more details.") - components.entry(frame, 7, 1, self.ui_state, "noising_bias", required=True) + self.components.label(frame, 7, 0, "Noising Bias", + tooltip="Controls the bias parameter of the timestep distribution function. Use the preview to see more details.") + self.components.entry(frame, 7, 1, ui_state, "noising_bias", required=True) # timestep shift - components.label(frame, 8, 0, "Timestep Shift", - tooltip="Shift the timestep distribution. Use the preview to see more details.") - components.entry(frame, 8, 1, self.ui_state, "timestep_shift", required=True) + self.components.label(frame, 8, 0, "Timestep Shift", + tooltip="Shift the timestep distribution. Use the preview to see more details.") + self.components.entry(frame, 8, 1, ui_state, "timestep_shift", required=True) if supports_dynamic_timestep_shifting: # dynamic timestep shifting - components.label(frame, 9, 0, "Dynamic Timestep Shifting", - tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) - components.switch(frame, 9, 1, self.ui_state, "dynamic_timestep_shifting") - - + self.components.label(frame, 9, 0, "Dynamic Timestep Shifting", + tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) + self.components.switch(frame, 9, 1, ui_state, "dynamic_timestep_shifting") - def __create_masked_frame(self, master, row): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) + def __create_masked_frame(self, master, row, ui_state): + frame = self.components.section_frame(master, row) # Masked Training - components.label(frame, 0, 0, "Masked Training", - tooltip="Masks the training samples to let the model focus on certain parts of the image. When enabled, one mask image is loaded for each training sample.") - components.switch(frame, 0, 1, self.ui_state, "masked_training") + self.components.label(frame, 0, 0, "Masked Training", + tooltip="Masks the training samples to let the model focus on certain parts of the image. When enabled, one mask image is loaded for each training sample.") + self.components.switch(frame, 0, 1, ui_state, "masked_training") # unmasked probability - components.label(frame, 1, 0, "Unmasked Probability", - tooltip="When masked training is enabled, specifies the number of training steps done on unmasked samples") - components.entry(frame, 1, 1, self.ui_state, "unmasked_probability", - extra_validate=check_range(lower=0, upper=1, message="Unmasked probability must be between 0 and 1")) + self.components.label(frame, 1, 0, "Unmasked Probability", + tooltip="When masked training is enabled, specifies the number of training steps done on unmasked samples") + self.components.entry(frame, 1, 1, ui_state, "unmasked_probability", + extra_validate=check_range(lower=0, upper=1, message="Unmasked probability must be between 0 and 1")) # unmasked weight - components.label(frame, 2, 0, "Unmasked Weight", - tooltip="When masked training is enabled, specifies the loss weight of areas outside the masked region") - components.entry(frame, 2, 1, self.ui_state, "unmasked_weight", - extra_validate=check_range(lower=0, upper=1, message="Unmasked weight must be between 0 and 1")) + self.components.label(frame, 2, 0, "Unmasked Weight", + tooltip="When masked training is enabled, specifies the loss weight of areas outside the masked region") + self.components.entry(frame, 2, 1, ui_state, "unmasked_weight", + extra_validate=check_range(lower=0, upper=1, message="Unmasked weight must be between 0 and 1")) # normalize masked area loss - components.label(frame, 3, 0, "Normalize Masked Area Loss", - tooltip="When masked training is enabled, normalizes the loss for each sample based on the sizes of the masked region") - components.switch(frame, 3, 1, self.ui_state, "normalize_masked_area_loss") + self.components.label(frame, 3, 0, "Normalize Masked Area Loss", + tooltip="When masked training is enabled, normalizes the loss for each sample based on the sizes of the masked region") + self.components.switch(frame, 3, 1, ui_state, "normalize_masked_area_loss") # masked prior preservation - components.label(frame, 4, 0, "Masked Prior Preservation Weight", - tooltip="Preserves regions outside the mask using the original untrained model output as a target. Only available for LoRA training. If enabled, use a low unmasked weight.") - components.entry(frame, 4, 1, self.ui_state, "masked_prior_preservation_weight", - extra_validate=check_range(lower=0, upper=1, message="Masked prior preservation weight must be between 0 and 1")) + self.components.label(frame, 4, 0, "Masked Prior Preservation Weight", + tooltip="Preserves regions outside the mask using the original untrained model output as a target. Only available for LoRA training. If enabled, use a low unmasked weight.") + self.components.entry(frame, 4, 1, ui_state, "masked_prior_preservation_weight", + extra_validate=check_range(lower=0, upper=1, message="Masked prior preservation weight must be between 0 and 1")) # use custom conditioning image - components.label(frame, 5, 0, "Custom Conditioning Image", - tooltip="When custom conditioning image is enabled, will use png postfix with -condlabel instead of automatically generated.It's suitable for special scenarios, such as object removal, allowing the model to learn a certain behavior concept") - components.switch(frame, 5, 1, self.ui_state, "custom_conditioning_image") + self.components.label(frame, 5, 0, "Custom Conditioning Image", + tooltip="When custom conditioning image is enabled, will use png postfix with -condlabel instead of automatically generated.It's suitable for special scenarios, such as object removal, allowing the model to learn a certain behavior concept") + self.components.switch(frame, 5, 1, ui_state, "custom_conditioning_image") - def __create_loss_frame(self, master, row, supports_vb_loss: bool = False): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) + def __create_loss_frame(self, master, row, controller, ui_state, + supports_vb_loss: bool = False): + frame = self.components.section_frame(master, row) # MSE Strength - components.label(frame, 0, 0, "MSE Strength", - tooltip="Mean Squared Error strength for custom loss settings. Strengths should generally sum to 1.") - components.entry(frame, 0, 1, self.ui_state, "mse_strength", required=True) + self.components.label(frame, 0, 0, "MSE Strength", + tooltip="Mean Squared Error strength for custom loss settings. Strengths should generally sum to 1.") + self.components.entry(frame, 0, 1, ui_state, "mse_strength", required=True) # MAE Strength - components.label(frame, 1, 0, "MAE Strength", - tooltip="Mean Absolute Error strength for custom loss settings. Strengths should generally sum to 1.") - components.entry(frame, 1, 1, self.ui_state, "mae_strength", required=True) + self.components.label(frame, 1, 0, "MAE Strength", + tooltip="Mean Absolute Error strength for custom loss settings. Strengths should generally sum to 1.") + self.components.entry(frame, 1, 1, ui_state, "mae_strength", required=True) # log-cosh Strength - components.label(frame, 2, 0, "log-cosh Strength", - tooltip="Log - Hyperbolic cosine Error strength for custom loss settings. Strengths should generally sum to 1.") - components.entry(frame, 2, 1, self.ui_state, "log_cosh_strength", required=True) + self.components.label(frame, 2, 0, "log-cosh Strength", + tooltip="Log - Hyperbolic cosine Error strength for custom loss settings. Strengths should generally sum to 1.") + self.components.entry(frame, 2, 1, ui_state, "log_cosh_strength", required=True) # Huber Strength - components.label(frame, 3, 0, "Huber Strength", - tooltip="Huber loss strength for custom loss settings. Less sensitive to outliers than MSE. Strengths should generally sum to 1.") - components.entry(frame, 3, 1, self.ui_state, "huber_strength", required=True) + self.components.label(frame, 3, 0, "Huber Strength", + tooltip="Huber loss strength for custom loss settings. Less sensitive to outliers than MSE. Strengths should generally sum to 1.") + self.components.entry(frame, 3, 1, ui_state, "huber_strength", required=True) # Huber Delta - components.label(frame, 4, 0, "Huber Delta", - tooltip="Delta parameter for huber loss") - components.entry(frame, 4, 1, self.ui_state, "huber_delta", required=True) + self.components.label(frame, 4, 0, "Huber Delta", + tooltip="Delta parameter for huber loss") + self.components.entry(frame, 4, 1, ui_state, "huber_delta", required=True) if supports_vb_loss: # VB Strength - components.label(frame, 5, 0, "VB Strength", - tooltip="Variational lower-bound strength for custom loss settings. Should be set to 1 for variational diffusion models") - components.entry(frame, 5, 1, self.ui_state, "vb_loss_strength", required=True) + self.components.label(frame, 5, 0, "VB Strength", + tooltip="Variational lower-bound strength for custom loss settings. Should be set to 1 for variational diffusion models") + self.components.entry(frame, 5, 1, ui_state, "vb_loss_strength", required=True) # Loss Weight function - components.label(frame, 6, 0, "Loss Weight Function", - tooltip="Choice of loss weight function. Can help the model learn details more accurately.") - components.options(frame, 6, 1, [str(x) for x in list(LossWeight) - if x.supports_flow_matching() == self.train_config.model_type.is_flow_matching() - or x == LossWeight.CONSTANT - ], - self.ui_state, "loss_weight_fn") + self.components.label(frame, 6, 0, "Loss Weight Function", + tooltip="Choice of loss weight function. Can help the model learn details more accurately.") + self.components.options(frame, 6, 1, [str(x) for x in list(LossWeight) + if x.supports_flow_matching() == controller.is_flow_matching() + or x == LossWeight.CONSTANT + ], + ui_state, "loss_weight_fn") row = 7 # Loss weight strength - if not self.train_config.model_type.is_flow_matching(): - components.label(frame, row, 0, "Gamma", - tooltip="Inverse strength of loss weighting. Range: 1-20, only applies to Min SNR and P2.") - components.entry(frame, row, 1, self.ui_state, "loss_weight_strength", - extra_validate=check_range(lower=1, upper=20, message="Gamma must be between 1 and 20")) + if not controller.is_flow_matching(): + self.components.label(frame, row, 0, "Gamma", + tooltip="Inverse strength of loss weighting. Range: 1-20, only applies to Min SNR and P2.") + self.components.entry(frame, row, 1, ui_state, "loss_weight_strength", + extra_validate=check_range(lower=1, upper=20, message="Gamma must be between 1 and 20")) row += 1 # Loss Scaler - components.label(frame, row, 0, "Loss Scaler", - tooltip="Selects the type of loss scaling to use during training. Functionally equated as: Loss * selection") - components.options(frame, row, 1, [str(x) for x in list(LossScaler)], self.ui_state, "loss_scaler") + self.components.label(frame, row, 0, "Loss Scaler", + tooltip="Selects the type of loss scaling to use during training. Functionally equated as: Loss * selection") + self.components.options(frame, row, 1, [str(x) for x in list(LossScaler)], ui_state, "loss_scaler") row += 1 - def __create_layer_frame(self, master, row): - cls = create.get_model_setup_class(self.train_config.model_type, self.train_config.training_method) - presets = cls.LAYER_PRESETS if cls is not None else {"full": []} - components.layer_filter_entry(master, row, 0, self.ui_state, - preset_var_name="layer_filter_preset", presets=presets, - preset_label="Layer Filter", - preset_tooltip="Select a preset defining which layers to train, or select 'Custom' to define your own.\nA blank 'custom' field or 'Full' will train all layers.", - entry_var_name="layer_filter", - entry_tooltip="Comma-separated list of diffusion layers to train. Regular expressions (if toggled) are supported. Any model layer with a matching name will be trained", - regex_var_name="layer_filter_regex", - regex_tooltip="If enabled, layer filter patterns are interpreted as regular expressions. Otherwise, simple substring matching is used.", - ) - - - def __on_layer_filter_preset_change(self): - if not self.layer_selector: - return - selected = self.ui_state.get_var("layer_filter_preset").get() - self.__preset_set_layer_choice(selected) - - def __hide_layer_entry(self): - if self.layer_entry and self.layer_entry.winfo_manager(): - self.layer_entry.grid_remove() - - def __show_layer_entry(self): - if self.layer_entry and not self.layer_entry.winfo_manager(): - self.layer_entry.grid() - - def __open_optimizer_params_window(self): - window = OptimizerParamsWindow(self.master, self.train_config, self.ui_state) - self.master.wait_window(window) - - def __open_scheduler_params_window(self): - window = SchedulerParamsWindow(self.master, self.train_config, self.ui_state) - self.master.wait_window(window) - - def __open_timestep_distribution_window(self): - window = TimestepDistributionWindow(self.master, self.train_config, self.ui_state) - self.master.wait_window(window) - - def __open_offloading_window(self): - window = OffloadingWindow(self.master, self.train_config, self.ui_state) - self.master.wait_window(window) - - def __restore_optimizer_config(self, *args): - optimizer_config = change_optimizer(self.train_config) - self.ui_state.get_var("optimizer").update(optimizer_config) - - def __restore_scheduler_config(self, variable): - if not hasattr(self, 'lr_scheduler_adv_comp'): - return - - if variable == "CUSTOM": - self.lr_scheduler_adv_comp.configure(state="normal") - else: - self.lr_scheduler_adv_comp.configure(state="disabled") + def __create_layer_frame(self, master, row, controller, ui_state): + presets = controller.get_layer_presets() + self.components.layer_filter_entry(master, row, 0, ui_state, + preset_var_name="layer_filter_preset", presets=presets, + preset_label="Layer Filter", + preset_tooltip="Select a preset defining which layers to train, or select 'Custom' to define your own.\nA blank 'custom' field or 'Full' will train all layers.", + entry_var_name="layer_filter", + entry_tooltip="Comma-separated list of diffusion layers to train. Regular expressions (if toggled) are supported. Any model layer with a matching name will be trained", + regex_var_name="layer_filter_regex", + regex_tooltip="If enabled, layer filter patterns are interpreted as regular expressions. Otherwise, simple substring matching is used.", + ) diff --git a/modules/ui/BaseVideoToolUIView.py b/modules/ui/BaseVideoToolUIView.py index c3291e6ea..b82ba5701 100644 --- a/modules/ui/BaseVideoToolUIView.py +++ b/modules/ui/BaseVideoToolUIView.py @@ -1,877 +1,196 @@ -import concurrent.futures -import math -import os -import pathlib -import random -import shlex -import subprocess -import threading import webbrowser -from fractions import Fraction -from tkinter import filedialog +from abc import ABC, abstractmethod -from modules.util.image_util import load_image from modules.util.path_util import SUPPORTED_VIDEO_EXTENSIONS -from modules.util.ui import components - -import av -import customtkinter as ctk -import cv2 -import scenedetect -from PIL import Image - - -class VideoToolUI(ctk.CTkToplevel): - def __init__( - self, - parent, - *args, **kwargs, - ): - ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) - - self.title("Video Tools") - self.geometry("600x720") - self.resizable(True, True) - self.wait_visibility() - self.focus_set() - - self.grid_rowconfigure(0, weight=1) - self.grid_rowconfigure(1, weight=0) - self.grid_columnconfigure(0, weight=1) - - tabview = ctk.CTkTabview(self) - tabview.grid(row=0, column=0, sticky="nsew") - - self.clip_extract_tab = self.__clip_extract_tab(tabview.add("extract clips")) - self.image_extract_tab = self.__image_extract_tab(tabview.add("extract images")) - self.video_download_tab = self.__video_download_tab(tabview.add("download")) - self.status_bar(self) - - def status_bar(self, master): - frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=1, column=0) - frame.grid_columnconfigure(0, weight=0, minsize=160) - frame.grid_columnconfigure(1, weight=0, minsize=300) - frame.grid_columnconfigure(2, weight=1) - - #create preview image - preview_path = "resources/icons/icon.png" - preview = load_image(preview_path, 'RGB') - preview.thumbnail((150, 150)) - self.preview_image= ctk.CTkImage(light_image=preview, size=preview.size) - self.preview_image_label = ctk.CTkLabel( - master=frame, text="Preview image", image=self.preview_image, height=150, width=150, - compound="top") - self.preview_image_label.grid(row=0, column=0, sticky="nw", padx=5, pady=5) - - #displays progress and messages that also go to terminal - self.status_label = ctk.CTkTextbox(master=frame, width=400, height=160, wrap="word", border_width=2) - self.status_label.insert(index="1.0", text="Current status") - self.status_label.configure(state="disabled") - self.status_label.grid(row=0, column=1, sticky="ne", padx=5, pady=5) - - def __clip_extract_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=120) - frame.grid_columnconfigure(1, weight=0, minsize=200) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) + +class BaseVideoToolUIView(ABC): + def __init__(self, components): + self.components = components + + def build_clip_extract_tab(self, frame, controller, ui_state): # single video - components.label(frame, 0, 0, "Single Video", + self.components.label(frame, 0, 0, "Single Video", tooltip="Link to single video file to process.") - self.clip_single_entry = ctk.CTkEntry(frame, width=190) - self.clip_single_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) - self.clip_single_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_file(self.clip_single_entry, - [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))] - )) - self.clip_single_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 0, 2, "Extract Single", - command=lambda: self.__extract_clips_button(False)) + self.components.entry(frame, 0, 1, ui_state, "clip_single", width=190) + self._create_browse_file_button(frame, 0, ui_state, "clip_single", + [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))]) + self.components.button(frame, 0, 2, "Extract Single", + command=lambda: self._extract_clips(False, controller)) # time range - components.label(frame, 1, 0, " Time Range", + self.components.label(frame, 1, 0, " Time Range", tooltip="Time range to limit selection for single video, \ format as hour:minute:second, minute:second, or seconds.") - self.clip_time_start_entry = ctk.CTkEntry(frame, width=100) - self.clip_time_start_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.clip_time_start_entry.insert(0, "00:00:00") - self.clip_time_end_entry = ctk.CTkEntry(frame, width=100) - self.clip_time_end_entry.grid(row=1, column=1, sticky="e", padx=5, pady=5) - self.clip_time_end_entry.insert(0, "99:99:99") + self.components.entry(frame, 1, 1, ui_state, "clip_time_start", width=100, sticky="w") + self.components.entry(frame, 1, 1, ui_state, "clip_time_end", width=100, sticky="e") # directory of videos - components.label(frame, 2, 0, "Directory", + self.components.label(frame, 2, 0, "Directory", tooltip="Path to directory with multiple videos to process, including in subdirectories.") - self.clip_list_entry = ctk.CTkEntry(frame, width=190) - self.clip_list_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - self.clip_list_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.clip_list_entry)) - self.clip_list_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 2, 2, "Extract Directory", - command=lambda: self.__extract_clips_button(True)) + self.components.entry(frame, 2, 1, ui_state, "clip_list", width=190) + self._create_browse_dir_button(frame, 2, ui_state, "clip_list") + self.components.button(frame, 2, 2, "Extract Directory", + command=lambda: self._extract_clips(True, controller)) # output directory - components.label(frame, 3, 0, "Output", + self.components.label(frame, 3, 0, "Output", tooltip="Path to folder where extracted clips will be saved.") - self.clip_output_entry = ctk.CTkEntry(frame, width=190) - self.clip_output_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) - self.clip_output_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.clip_output_entry)) - self.clip_output_button.grid(row=3, column=1, sticky="e", padx=5, pady=5) + self.components.entry(frame, 3, 1, ui_state, "clip_output", width=190) + self._create_browse_dir_button(frame, 3, ui_state, "clip_output") # output to subdirectories - self.output_subdir_clip = ctk.BooleanVar(self, False) - components.label(frame, 4, 0, "Output to\nSubdirectories", + self.components.label(frame, 4, 0, "Output to\nSubdirectories", tooltip="If enabled, files are saved to subfolders based on filename and input directory. \ Otherwise will all be saved to the top level of the output directory.") - self.output_subdir_clip_entry = ctk.CTkSwitch(frame, variable=self.output_subdir_clip, text="") - self.output_subdir_clip_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + self.components.switch(frame, 4, 1, ui_state, "output_subdir_clip") # split at cuts - self.split_at_cuts = ctk.BooleanVar(self, False) - components.label(frame, 5, 0, "Split at Cuts", + self.components.label(frame, 5, 0, "Split at Cuts", tooltip="If enabled, detect cuts in the input video and split at those points. \ Otherwise will split at any point, and clips may contain cuts.") - self.split_cuts_entry = ctk.CTkSwitch(frame, variable=self.split_at_cuts, text="") - self.split_cuts_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) + self.components.switch(frame, 5, 1, ui_state, "split_cuts") # maximum length - components.label(frame, 6, 0, "Max Length (s)", + self.components.label(frame, 6, 0, "Max Length (s)", tooltip="Maximum length in seconds for saved clips, larger clips will be broken into multiple small clips.") - self.clip_length_entry = ctk.CTkEntry(frame, width=220) - self.clip_length_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) - self.clip_length_entry.insert(0, "3") + self.components.entry(frame, 6, 1, ui_state, "clip_length", width=220) # Set FPS - components.label(frame, 7, 0, "Set FPS", + self.components.label(frame, 7, 0, "Set FPS", tooltip="FPS to convert output videos to, set to 0 to keep original rate.") - self.clip_fps_entry = ctk.CTkEntry(frame, width=220) - self.clip_fps_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) - self.clip_fps_entry.insert(0, "24.0") + self.components.entry(frame, 7, 1, ui_state, "clip_fps", width=220) # Remove borders - self.clip_bordercrop = ctk.BooleanVar(self, False) - components.label(frame, 8, 0, "Remove Borders", + self.components.label(frame, 8, 0, "Remove Borders", tooltip="Remove black borders from output clip") - self.clip_bordercrop_entry = ctk.CTkSwitch(frame, variable=self.clip_bordercrop, text="") - self.clip_bordercrop_entry.grid(row=8, column=1, sticky="w", padx=5, pady=5) + self.components.switch(frame, 8, 1, ui_state, "clip_bordercrop") # Crop Variation - components.label(frame, 9, 0, "Crop Variation", + self.components.label(frame, 9, 0, "Crop Variation", tooltip="Output clips will be randomly cropped to +- the base aspect ratio, \ somewhat biased towards making square videos. Set to 0 to use only base aspect.") - self.clip_crop_entry = ctk.CTkEntry(frame, width=220) - self.clip_crop_entry.grid(row=9, column=1, sticky="w", padx=5, pady=5) - self.clip_crop_entry.insert(0, "0.2") - - # object filter - currently unused, may implement in future - # components.label(frame, 9, 0, "Object Filter", - # tooltip="Detect general features using Haar-Cascade classifier, and choose how to deal with clips where it is detected") - # components.options(frame, 9, 1, ["NONE", "FACE", "EYE", "BODY"], self.video_ui_state, "filter_object") - # components.options(frame, 9, 2, ["INCLUDE", "EXCLUDE", "SUBFOLDER"], self.video_ui_state, "filter_behavior") - - frame.pack(fill="both", expand=1) - return frame - - def __image_extract_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=120) - frame.grid_columnconfigure(1, weight=0, minsize=200) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) + self.components.entry(frame, 9, 1, ui_state, "clip_crop", width=220) + def build_image_extract_tab(self, frame, controller, ui_state): # single video - components.label(frame, 0, 0, "Single Video", + self.components.label(frame, 0, 0, "Single Video", tooltip="Link to single video file to process.") - self.image_single_entry = ctk.CTkEntry(frame, width=190) - self.image_single_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) - self.image_single_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_file(self.image_single_entry, - [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))] - )) - self.image_single_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 0, 2, "Extract Single", - command=lambda: self.__extract_images_button(False)) + self.components.entry(frame, 0, 1, ui_state, "image_single", width=190) + self._create_browse_file_button(frame, 0, ui_state, "image_single", + [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))]) + self.components.button(frame, 0, 2, "Extract Single", + command=lambda: self._extract_images(False, controller)) # time range - components.label(frame, 1, 0, " Time Range", + self.components.label(frame, 1, 0, " Time Range", tooltip="Time range to limit selection for single video, \ format as hour:minute:second, minute:second, or seconds.") - self.image_time_start_entry = ctk.CTkEntry(frame, width=100) - self.image_time_start_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.image_time_start_entry.insert(0, "00:00:00") - self.image_time_end_entry = ctk.CTkEntry(frame, width=100) - self.image_time_end_entry.grid(row=1, column=1, sticky="e", padx=5, pady=5) - self.image_time_end_entry.insert(0, "99:99:99") + self.components.entry(frame, 1, 1, ui_state, "image_time_start", width=100, sticky="w") + self.components.entry(frame, 1, 1, ui_state, "image_time_end", width=100, sticky="e") # directory of videos - components.label(frame, 2, 0, "Directory", + self.components.label(frame, 2, 0, "Directory", tooltip="Path to directory with multiple videos to process, including in subdirectories.") - self.image_list_entry = ctk.CTkEntry(frame, width=190) - self.image_list_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - self.image_list_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.image_list_entry)) - self.image_list_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 2, 2, "Extract Directory", - command=lambda: self.__extract_images_button(True)) + self.components.entry(frame, 2, 1, ui_state, "image_list", width=190) + self._create_browse_dir_button(frame, 2, ui_state, "image_list") + self.components.button(frame, 2, 2, "Extract Directory", + command=lambda: self._extract_images(True, controller)) # output directory - components.label(frame, 3, 0, "Output", + self.components.label(frame, 3, 0, "Output", tooltip="Path to folder where extracted images will be saved.") - self.image_output_entry = ctk.CTkEntry(frame, width=190) - self.image_output_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) - self.image_output_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.image_output_entry)) - self.image_output_button.grid(row=3, column=1, sticky="e", padx=5, pady=5) + self.components.entry(frame, 3, 1, ui_state, "image_output", width=190) + self._create_browse_dir_button(frame, 3, ui_state, "image_output") # output to subdirectories - self.output_subdir_img = ctk.BooleanVar(self, False) - components.label(frame, 4, 0, "Output to\nSubdirectories", + self.components.label(frame, 4, 0, "Output to\nSubdirectories", tooltip="If enabled, files are saved to subfolders based on filename and input directory. \ Otherwise will all be saved to the top level of the output directory.") - self.output_subdir_img_entry = ctk.CTkSwitch(frame, variable=self.output_subdir_img, text="") - self.output_subdir_img_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + self.components.switch(frame, 4, 1, ui_state, "output_subdir_img") # image capture rate - components.label(frame, 5, 0, "Images/sec", + self.components.label(frame, 5, 0, "Images/sec", tooltip="Number of images to capture per second of video. \ Images will be taken at semi-random frames around the specified frequency.") - self.capture_rate_entry = ctk.CTkEntry(frame, width=220) - self.capture_rate_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) - self.capture_rate_entry.insert(0, "0.5") + self.components.entry(frame, 5, 1, ui_state, "capture_rate", width=220) # blur removal - components.label(frame, 6, 0, "Blur Removal", + self.components.label(frame, 6, 0, "Blur Removal", tooltip="Threshold for removal of blurry images, relative to all others. \ For example at 0.2, the blurriest 20%% of the final selected frames will not be saved.") - self.blur_threshold_entry = ctk.CTkEntry(frame, width=220) - self.blur_threshold_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) - self.blur_threshold_entry.insert(0, "0.2") + self.components.entry(frame, 6, 1, ui_state, "blur_threshold", width=220) # Remove borders - self.image_bordercrop = ctk.BooleanVar(self, False) - components.label(frame, 7, 0, "Remove Borders", + self.components.label(frame, 7, 0, "Remove Borders", tooltip="Remove black borders from output image") - self.image_bordercrop_entry = ctk.CTkSwitch(frame, variable=self.image_bordercrop, text="") - self.image_bordercrop_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) + self.components.switch(frame, 7, 1, ui_state, "image_bordercrop") # Crop Variation - components.label(frame, 8, 0, "Crop Variation", + self.components.label(frame, 8, 0, "Crop Variation", tooltip="Output images will be randomly cropped to +- the base aspect ratio, \ somewhat biased towards making square images. Set to 0 to use only base sapect.") - self.image_crop_entry = ctk.CTkEntry(frame, width=220) - self.image_crop_entry.grid(row=8, column=1, sticky="w", padx=5, pady=5) - self.image_crop_entry.insert(0, "0.2") - - # # object filter - currently unused, may implement in future - # components.label(frame, 5, 0, "Object Filter", - # tooltip="Detect general features using Haar-Cascade classifier, and choose how to deal with clips where it is detected") - # components.options(frame, 5, 1, ["NONE", "FACE", "EYE", "BODY"], self.video_ui_state, "filter_object") - # components.options(frame, 5, 2, ["INCLUDE", "EXCLUDE", "SUBFOLDER"], self.video_ui_state, "filter_behavior") - - frame.pack(fill="both", expand=1) - return frame - - def __video_download_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=120) - frame.grid_columnconfigure(1, weight=0, minsize=200) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) + self.components.entry(frame, 8, 1, ui_state, "image_crop", width=220) + def build_video_download_tab(self, frame, controller, ui_state): # link - components.label(frame, 0, 0, "Single Link", + self.components.label(frame, 0, 0, "Single Link", tooltip="Link to video/playlist to download. Uses yt-dlp, supports youtube, twitch, instagram, and many other sites.") - self.download_link_entry = ctk.CTkEntry(frame, width=220) - self.download_link_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) - components.button(frame, 0, 2, "Download Link", command=lambda: self.__download_button(False)) + self.components.entry(frame, 0, 1, ui_state, "download_link", width=220) + self.components.button(frame, 0, 2, "Download Link", + command=lambda: self._download(False, controller)) # link list - components.label(frame, 1, 0, "Link List", + self.components.label(frame, 1, 0, "Link List", tooltip="Path to txt file with list of links separated by newlines.") - self.download_list_entry = ctk.CTkEntry(frame, width=190) - self.download_list_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.download_list_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_file(self.download_list_entry, [("Text file", ".txt")])) - self.download_list_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 1, 2, "Download List", command=lambda: self.__download_button(True)) + self.components.entry(frame, 1, 1, ui_state, "download_list", width=190) + self._create_browse_file_button(frame, 1, ui_state, "download_list", [("Text file", ".txt")]) + self.components.button(frame, 1, 2, "Download List", + command=lambda: self._download(True, controller)) # output directory - components.label(frame, 2, 0, "Output", + self.components.label(frame, 2, 0, "Output", tooltip="Path to folder where downloaded videos will be saved.") - self.download_output_entry = ctk.CTkEntry(frame, width=190) - self.download_output_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - self.download_output_button = ctk.CTkButton(frame, width=30, text="...", command=lambda: self.__browse_for_dir(self.download_output_entry)) - self.download_output_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) + self.components.entry(frame, 2, 1, ui_state, "download_output", width=190) + self._create_browse_dir_button(frame, 2, ui_state, "download_output") # additional args - components.label(frame, 3, 0, "Additional Args", + self.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.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 --format mp4") - components.button(frame, 3, 2, "yt-dlp info", + self._create_textbox(frame, 3, 1, 220, 90, ui_state, "download_args") + self.components.button(frame, 3, 2, "yt-dlp info", command=lambda: webbrowser.open("https://github.com/yt-dlp/yt-dlp?tab=readme-ov-file#usage-and-options", new=0, autoraise=False)) - frame.pack(fill="both", expand=1) - return frame - - def __browse_for_dir(self, entry_box): - # get the path from the user - path = filedialog.askdirectory() - # set the path to the entry box - # delete entry box text - entry_box.focus_set() - entry_box.delete(0, ctk.END) - entry_box.insert(0, path) - self.focus_set() - - def __browse_for_file(self, entry_box, filetypes): - # get the path from the user - path = filedialog.askopenfilename(filetypes=filetypes) - # set the path to the entry box - # delete entry box text - entry_box.focus_set() - entry_box.delete(0, ctk.END) - entry_box.insert(0, path) - self.focus_set() - - def __get_vid_paths(self, batch_mode: bool, input_path_single: str, input_path_dir: str): - input_videos = [] - if not batch_mode: - path = pathlib.Path(input_path_single) - if path.is_file(): - vid = cv2.VideoCapture(str(path)) - ok = False - try: - if vid.isOpened(): - ok, _ = vid.read() - finally: - vid.release() - if ok: - return [path] - else: - self.__update_status("Invalid video file!") - return [] - else: - self.__update_status("No file specified, or invalid file path!") - return [] - else: - input_videos = [] - if not pathlib.Path(input_path_dir).is_dir() or input_path_dir == "": - self.__update_status("Invalid input directory!") - return [] - # Only traverse supported extensions to avoid opening every file. - lower_exts = {e.lower() for e in SUPPORTED_VIDEO_EXTENSIONS} - for path in pathlib.Path(input_path_dir).rglob("*"): - if path.is_file() and path.suffix.lower() in lower_exts: - vid = cv2.VideoCapture(str(path)) - ok = False - try: - if vid.isOpened(): - ok, _ = vid.read() - finally: - vid.release() - if ok: - input_videos.append(path) - self.__update_status(f'Found {len(input_videos)} videos to process') - return input_videos - - def __run_in_thread(self, target, *args): - """Clear status box and run target function in a daemon thread.""" - self.status_label.configure(state="normal") - self.status_label.delete(index1="1.0", index2="end") - self.status_label.configure(state="disabled") - t = threading.Thread(target=target, args=args) - t.daemon = True - t.start() - - @staticmethod - def __parse_timestamp_to_frames(timestamp: str, fps: float) -> int: - return int(sum(int(x) * 60 ** i for i, x in enumerate(reversed(timestamp.split(':')))) * fps) - - def __get_safe_fps(self, video: cv2.VideoCapture, video_path: str) -> float: - fps = video.get(cv2.CAP_PROP_FPS) or 0.0 - if fps <= 0: - self.__update_status(f'Warning: Could not read FPS for "{os.path.basename(video_path)}". Falling back to 30 FPS.') - return 30.0 - return fps - - @staticmethod - def __get_output_dir(use_subdir: bool, batch_mode: bool, output_entry: str, - video_path, input_dir: str) -> str: - if use_subdir and batch_mode: - return os.path.join(output_entry, - os.path.splitext(os.path.relpath(video_path, input_dir))[0]) - elif use_subdir: - return os.path.join(output_entry, - os.path.splitext(os.path.basename(video_path))[0]) - return output_entry - - def __get_random_aspect(self, height: int, width: int, variation: float) -> tuple[int, int, int, int]: - # Return original dimensions and no offset if variation is zero - if variation == 0: - return 0, height, 0, width - - old_aspect = height/width - variation_scaled = old_aspect*variation - if old_aspect > 1.2: #tall image - new_aspect = min(4.0, max(1.0, random.triangular(old_aspect-(variation_scaled*1.5), old_aspect+(variation_scaled/2), old_aspect))) - elif old_aspect < 0.85: #wide image - new_aspect = max(0.25, min(1.0, random.triangular(old_aspect-(variation_scaled/2), old_aspect+(variation_scaled*1.5), old_aspect))) - else: #square image - new_aspect = random.triangular(old_aspect-variation_scaled, old_aspect+variation_scaled) - - new_aspect = round(new_aspect, 2) - #keep the height the same if reducing width, and vice versa - if new_aspect > old_aspect: - new_height = int(height) - new_width = int(width*(old_aspect/new_aspect)) - elif new_aspect < old_aspect: - new_height = int(height*(new_aspect/old_aspect)) - new_width = int(width) - else: - new_height = int(height) - new_width = int(width) - - #random offset in dimension that was cropped - position_x = random.randint(0, width-new_width) - position_y = random.randint(0, height-new_height) - return position_y, new_height, position_x, new_width - - def find_main_contour(self, frame): - #outline image to find main content and exclude black bars often present on letterboxed videos - frame_grayscale = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) - _, frame_thresh = cv2.threshold(frame_grayscale, 15, 255, cv2.THRESH_BINARY) - frame_contours, _ = cv2.findContours(frame_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) - if frame_contours: - #select largest contour by area - frame_maincontour = max(frame_contours, key=lambda c: cv2.contourArea(c)) - x1, y1, w1, h1 = cv2.boundingRect(frame_maincontour) - else: #fallback if no contours detected - x1 = 0 - y1 = 0 - h1, w1, _ = frame.shape - - #if bounding box did not detect the correct area, likely due to all-black frame - if not frame_contours or h1 < 10 or w1 < 10: - x1 = 0 - y1 = 0 - h1, w1, _ = frame.shape - return x1, y1, w1, h1 - - def __extract_clips_button(self, batch_mode: bool): - self.__run_in_thread(self.__extract_clips_multi, batch_mode) - - def __extract_clips_multi(self, batch_mode: bool): - if not pathlib.Path(self.clip_output_entry.get()).is_dir() or self.clip_output_entry.get() == "": - self.__update_status("Invalid output directory!") - return - - # validate numeric inputs - try: - max_length = float(self.clip_length_entry.get()) - crop_variation = float(self.clip_crop_entry.get()) - target_fps = float(self.clip_fps_entry.get()) - input_single_entry = self.clip_single_entry.get() - input_multiple_entry = self.clip_list_entry.get() - output_entry = self.clip_output_entry.get() - except ValueError: - self.__update_status("Invalid numeric input for Max Length, Crop Variation, or FPS.") - return - if max_length <= 0.25: - self.__update_status("Max Length of clips must be > 0.25 seconds.") - return - if target_fps < 0: - self.__update_status("Target FPS must be a positive number (or 0 to skip fps re-encoding).") - return - if not (0.0 <= crop_variation < 1.0): - self.__update_status("Crop Variation must be between 0.0 and 1.0.") - return - - input_videos = self.__get_vid_paths(batch_mode, input_single_entry, input_multiple_entry) - if len(input_videos) == 0: # exit if no paths found - return - - with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: - for video_path in input_videos: - output_directory = self.__get_output_dir( - self.output_subdir_clip_entry.get(), batch_mode, - output_entry, video_path, input_multiple_entry) - time_start = "00:00:00" if batch_mode else str(self.clip_time_start_entry.get()) - time_end = "99:99:99" if batch_mode else str(self.clip_time_end_entry.get()) - executor.submit(self.__extract_clips, - str(video_path), time_start, time_end, max_length, - self.split_at_cuts.get(), bool(self.clip_bordercrop_entry.get()), - crop_variation, target_fps, output_directory) - - if batch_mode: - self.__update_status(f'Clip extraction from all videos in "{input_multiple_entry}" complete') - else: - self.__update_status(f'Clip extraction from "{input_single_entry}" complete') - - def __extract_clips(self, video_path: str, timestamp_min: str, timestamp_max: str, max_length: float, - split_at_cuts: bool, remove_borders: bool, crop_variation: float, target_fps: float, output_dir: str): - video = cv2.VideoCapture(video_path) - vid_fps = self.__get_safe_fps(video, video_path) - max_length_frames = int(max_length * vid_fps) - min_length_frames = max(int(0.25 * vid_fps), 1) - total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) or 0 - timestamp_max_frame = min(self.__parse_timestamp_to_frames(timestamp_max, vid_fps), max(total_frames - 1, 0)) - timestamp_min_frame = min(self.__parse_timestamp_to_frames(timestamp_min, vid_fps), timestamp_max_frame) - - if split_at_cuts: - #use scenedetect to find cuts, based on start/end frame number - self.__update_status(f'Detecting scenes in "{os.path.basename(video_path)}"') - timecode_list = scenedetect.detect( - video_path=str(video_path), - detector=scenedetect.AdaptiveDetector(), - start_time=int(timestamp_min_frame), - end_time=int(timestamp_max_frame)) - scene_list = [(x[0].get_frames(), x[1].get_frames()) for x in timecode_list] - if not scene_list: - scene_list = [(timestamp_min_frame, timestamp_max_frame)] - else: - scene_list = [(timestamp_min_frame, timestamp_max_frame)] - - scene_list_split = [] - for scene in scene_list: - length = scene[1]-scene[0] - if length > max_length_frames: #check for any scenes longer than max length - n = math.ceil(length/max_length_frames) #divide into n new scenes - new_length = int(length/n) - new_splits = range(scene[0], scene[1]+min_length_frames, new_length) #divide clip into closest chunks to max_length - for i, _n in enumerate(new_splits[:-1]): - if new_splits[i + 1] - new_splits[i] > min_length_frames: - scene_list_split.append((new_splits[i], new_splits[i + 1])) - elif length > (min_length_frames + 2): - # Trim first/last frame to avoid transition artifacts - scene_list_split.append((scene[0] + 1, scene[1] - 1)) - - self.__update_status(f'Video "{os.path.basename(video_path)}" being split into {len(scene_list_split)} clips in "{output_dir}"') - - with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: - futures = [ - executor.submit(self.__save_clip, scene, video_path, target_fps, - remove_borders, crop_variation, output_dir) - for scene in scene_list_split - ] - for future in concurrent.futures.as_completed(futures): - exc = future.exception() - if exc is not None: - self.__update_status(f'Error saving clip: {exc}') - - video.release() - - def __save_clip(self, scene: tuple[int, int], video_path: str, target_fps: float, - remove_borders: bool, crop_variation: float, output_dir: str): - basename, ext = os.path.splitext(os.path.basename(video_path)) - video = cv2.VideoCapture(str(video_path)) - fps = self.__get_safe_fps(video, video_path) - os.makedirs(output_dir, exist_ok=True) - output_name = f'{output_dir}{os.sep}{basename}_{scene[0]}-{scene[1]}' - output_ext = ".mp4" - - video.set(cv2.CAP_PROP_POS_FRAMES, (scene[1] + scene[0])//2) #set to middle of scene - frame_number = int(video.get(cv2.CAP_PROP_POS_FRAMES)) - success, frame = video.read() - if not success or frame is None: - self.__update_status(f'Failed to read frame from "{os.path.basename(video_path)}" at {int(frame_number)}. Skipping clip.') - video.release() - return - - # Blend random frames to detect borders, avoiding incorrect crop from black frames - if remove_borders: - frame_blend = frame - for i in range(5): - random_frame = random.randint(scene[0], scene[1]) - video.set(cv2.CAP_PROP_POS_FRAMES, random_frame) - success, frame = video.read() - if not success or frame is None: - continue - a = 1/(i+1) - b = 1-a - frame_blend = cv2.addWeighted(frame, a, frame_blend, b, 0) - x1, y1, w1, h1 = self.find_main_contour(frame_blend) - else: - x1 = 0 - y1 = 0 - h1, w1, _ = frame.shape - - y2, h2, x2, w2 = self.__get_random_aspect(h1, w1, crop_variation) - # Ensure dimensions are even, required - h2 -= h2 % 2 - w2 -= w2 % 2 - print(end='\x1b[2K') #clear terminal so next line can overwrite it - print(f'Saving frames {scene[0]}-{scene[1]} at size {w2}x{h2}', end="\r") - video.set(cv2.CAP_PROP_POS_FRAMES, (scene[1] + scene[0])//2) - success, frame = video.read() - if success: - try: - preview = Image.fromarray( - cv2.cvtColor(frame[y1+y2:y1+y2+h2, x1+x2:x1+x2+w2], cv2.COLOR_BGR2RGB)) - preview.thumbnail((150, 150)) - self.preview_image.configure(light_image=preview, size=preview.size) - #truncate filename of long files so UI doesn't shift around - filename_truncated = basename + ext if len(basename) < 20 else basename[:18] + ".." + ext - self.preview_image_label.configure( - text=f'{filename_truncated}\nFrames: {scene[0]}-{scene[1]}\nSize: {w2}x{h2}') - except Exception: - pass - video.release() - - if target_fps <= 0: - target_fps = fps - - output_path = f'{output_name}{output_ext}' - self.__write_clip_av(video_path, output_path, scene, fps, target_fps, - x1 + x2, y1 + y2, w2, h2) - - @staticmethod - def __write_clip_av(video_path: str, output_path: str, scene: tuple[int, int], - src_fps: float, target_fps: float, - crop_x: int, crop_y: int, crop_w: int, crop_h: int): - start_sec = scene[0] / src_fps - end_sec = scene[1] / src_fps - rate_frac = Fraction(target_fps).limit_denominator(10000) - stream_time_base = Fraction(rate_frac.denominator, rate_frac.numerator) - - with av.open(video_path) as input_container: - in_video = input_container.streams.video[0] - in_video.thread_type = 'AUTO' - in_audio = input_container.streams.audio[0] if input_container.streams.audio else None - - with av.open(output_path, mode='w') as output_container: - out_video = output_container.add_stream('libx264', rate=rate_frac) - out_video.width = crop_w - out_video.height = crop_h - out_video.pix_fmt = 'yuv420p' - out_video.time_base = stream_time_base - - out_audio = output_container.add_stream_from_template(in_audio) if in_audio else None - - input_container.seek(int(start_sec * 1_000_000)) - - out_frame_idx = 0 - out_time_step = 1.0 / target_fps - video_done = False - decode_streams = [s for s in (in_video, in_audio) if s is not None] - - for packet in input_container.demux(decode_streams): - if packet.stream == in_video: - if video_done: - continue - for frame in packet.decode(): - if frame.time is None or frame.time < start_sec: - continue - if frame.time >= end_sec: - video_done = True - break - - # FPS conversion: skip frames when source fps > target fps - if frame.time < start_sec + out_frame_idx * out_time_step: - continue - - img = frame.to_ndarray(format='bgr24') - cropped = img[crop_y:crop_y + crop_h, crop_x:crop_x + crop_w] - out_frame = av.VideoFrame.from_ndarray(cropped, format='bgr24') - out_frame.pts = out_frame_idx - out_frame.time_base = stream_time_base - - for out_pkt in out_video.encode(out_frame): - output_container.mux(out_pkt) - out_frame_idx += 1 - - elif packet.stream == in_audio and out_audio is not None: - if packet.dts is None: - continue - pkt_time = float(packet.pts * packet.time_base) - if pkt_time < start_sec or pkt_time >= end_sec: - continue - # Re-timestamp audio relative to clip start - packet.pts = int((pkt_time - start_sec) / packet.time_base) - packet.dts = packet.pts - packet.stream = out_audio - output_container.mux(packet) - - # Flush video encoder - for pkt in out_video.encode(): - output_container.mux(pkt) - - def __extract_images_button(self, batch_mode: bool): - self.__run_in_thread(self.__extract_images_multi, batch_mode) - - def __extract_images_multi(self, batch_mode : bool): - if not pathlib.Path(self.image_output_entry.get()).is_dir() or self.image_output_entry.get() == "": - self.__update_status("Invalid output directory!") - return - - # validate numeric inputs - try: - capture_rate = float(self.capture_rate_entry.get()) - blur_threshold = float(self.blur_threshold_entry.get()) - crop_variation = float(self.image_crop_entry.get()) - input_single_entry = self.image_single_entry.get() - input_multiple_entry = self.image_list_entry.get() - output_entry = self.image_output_entry.get() - except ValueError: - self.__update_status("Invalid numeric input for Images/sec, Blur Removal, or Crop Variation.") - return - if capture_rate <= 0: - self.__update_status("Images/sec must be > 0.") - return - if not (0.0 <= blur_threshold < 1.0): - self.__update_status("Blur Removal must be between 0.0 and 1.0.") - return - if not (0.0 <= crop_variation < 1.0): - self.__update_status("Crop Variation must be between 0.0 and 1.0.") - return - - input_videos = self.__get_vid_paths(batch_mode, input_single_entry, input_multiple_entry) - if not input_videos: - return - - with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: - for video_path in input_videos: - output_directory = self.__get_output_dir( - self.output_subdir_img_entry.get(), batch_mode, - output_entry, video_path, input_multiple_entry) - time_start = "00:00:00" if batch_mode else str(self.image_time_start_entry.get()) - time_end = "99:99:99" if batch_mode else str(self.image_time_end_entry.get()) - executor.submit(self.__save_frames, - str(video_path), time_start, time_end, capture_rate, - blur_threshold, self.image_bordercrop.get(), - crop_variation, output_directory) - if batch_mode: - self.__update_status(f'Image extraction from all videos in {input_multiple_entry} complete') - else: - self.__update_status(f'Image extraction from "{input_single_entry}" complete') - - def __save_frames(self, video_path: str, timestamp_min: str, timestamp_max: str, capture_rate: float, - blur_threshold: float, remove_borders: bool, crop_variation: float, output_dir: str): - video = cv2.VideoCapture(video_path) - vid_fps = self.__get_safe_fps(video, video_path) - if capture_rate <= 0: - self.__update_status("Images/sec must be > 0.") - video.release() - return - image_rate = max(int(vid_fps / capture_rate), 1) - total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) or 0 - timestamp_max_frame = min(self.__parse_timestamp_to_frames(timestamp_max, vid_fps), max(total_frames - 1, 0)) - timestamp_min_frame = min(self.__parse_timestamp_to_frames(timestamp_min, vid_fps), timestamp_max_frame) - frame_range = range(timestamp_min_frame, timestamp_max_frame, image_rate) - frame_list = [] - - for n in frame_range: - #pick frame from random triangular distribution around center of each "chunk" of the video - frame = abs(int(random.triangular(n-(image_rate/2), n+(image_rate/2)))) - frame = max(0, min(frame, max(total_frames - 1, 0))) - frame_list.append(frame) - - self.__update_status(f'Video "{os.path.basename(video_path)}" will be split into {len(frame_list)} images in "{output_dir}"') - - output_list = [] - for f in frame_list: - video.set(cv2.CAP_PROP_POS_FRAMES, f) - success, frame = video.read() - if success and frame is not None: - frame_grayscale = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) - frame_sharpness = cv2.Laplacian(frame_grayscale, cv2.CV_64F).var() - output_list.append((f, frame_sharpness)) - - if not output_list: - self.__update_status(f'No frames extracted from "{os.path.basename(video_path)}" in the selected range.') - video.release() - return - - output_list_sorted = sorted(output_list, key=lambda x: x[1]) - cutoff = int(blur_threshold * len(output_list_sorted)) - output_list_cut = output_list_sorted[cutoff:] - self.__update_status(f'{cutoff} blurriest images have been dropped from "{os.path.basename(video_path)}"') - - basename, ext = os.path.splitext(os.path.basename(video_path)) - os.makedirs(output_dir, exist_ok=True) - - for f in output_list_cut: - filename = f'{output_dir}{os.sep}{basename}_{f[0]}.jpg' - video.set(cv2.CAP_PROP_POS_FRAMES, f[0]) - success, frame = video.read() - - #crop out borders of frame - if remove_borders and success and frame is not None: - x1, y1, w1, h1 = self.find_main_contour(frame) - frame_cropped = frame[y1:y1+h1, x1:x1+w1] - else: - frame_cropped = frame if success and frame is not None else None - if frame_cropped is not None: - x1 = 0 - y1 = 0 - h1, w1, _ = frame_cropped.shape - - y2, h2, x2, w2 = self.__get_random_aspect(h1, w1, crop_variation) - - if success and frame is not None and frame_cropped is not None: - print(end='\x1b[2K') #clear terminal so next line can overwrite it - print(f'Saving frame {f[0]} at size {w2}x{h2}', end="\r") - try: - preview = Image.fromarray( - cv2.cvtColor(frame_cropped[y2:y2+h2, x2:x2+w2], cv2.COLOR_BGR2RGB)) - preview.thumbnail((150, 150)) - filename_truncated = basename + ext if len(basename) < 20 else basename[:17] + "..." + ext - self.preview_image.configure(light_image=preview, size=preview.size) - self.preview_image_label.configure(text=f'{filename_truncated}\nFrame: {f[0]}\nSize: {w2}x{h2}') - except Exception: - pass # preview update is non-critical - - cv2.imwrite(filename, frame_cropped[y2:y2+h2, x2:x2+w2]) - video.release() - - def __download_button(self, batch_mode: bool): - self.__run_in_thread(self.__download_multi, batch_mode) - - def __update_status(self, status_text: str): - print(status_text) - self.status_label.configure(state="normal") - self.status_label.insert(index="end", text=status_text + "\n") - self.status_label.configure(state="disabled") - - def __download_multi(self, batch_mode: bool): - if not pathlib.Path(self.download_output_entry.get()).is_dir() or self.download_output_entry.get() == "": - self.__update_status("Invalid output directory!") - return - - if not batch_mode: - ydl_urls = [self.download_link_entry.get()] - elif batch_mode: - ydl_path = pathlib.Path(self.download_list_entry.get()) - if ydl_path.is_file() and ydl_path.suffix.lower() == ".txt": - with open(ydl_path) as file: - ydl_urls = file.readlines() - else: - self.__update_status("Invalid link list!") - return - - with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: - for url in ydl_urls: - executor.submit(self.__download_video, - url.strip(), self.download_output_entry.get(), - self.download_args_entry.get("0.0", ctk.END)) - - self.__update_status(f'Completed {len(ydl_urls)} downloads.') - - def __download_video(self, url: str, output_dir: str, output_args: str): - url = (url or "").strip() - if not url: - self.__update_status("Empty URL, skipping download.") - return - - #Respect quotes and split into list to run as yt-dlp command - additional_args = shlex.split(output_args.strip()) if output_args and output_args.strip() else [] - cmd = ["yt-dlp", "-o", "%(title)s.%(ext)s", "-P", output_dir] + additional_args + [url] - - self.__update_status(f'Downloading {url}') - subprocess.run(cmd) - self.__update_status(f'Download {url} done!') + @abstractmethod + def _create_textbox(self, master, row, col, width, height, ui_state, var_name): + pass + + @abstractmethod + def _create_browse_dir_button(self, master, row, ui_state, var_name): + pass + + @abstractmethod + def _create_browse_file_button(self, master, row, ui_state, var_name, filetypes): + pass + + @abstractmethod + def update_status(self, status_text: str): + pass + + @abstractmethod + def clear_status(self): + pass + + @abstractmethod + def update_preview(self, preview_image, label_text: str): + pass + + def _extract_clips(self, batch_mode: bool, controller): + controller.extract_clips_button(batch_mode) + + def _extract_images(self, batch_mode: bool, controller): + controller.extract_images_button(batch_mode) + + def _download(self, batch_mode: bool, controller): + controller.download_button(batch_mode) diff --git a/modules/ui/CaptionUIController.py b/modules/ui/CaptionUIController.py index e6cc0551e..0495322a3 100644 --- a/modules/ui/CaptionUIController.py +++ b/modules/ui/CaptionUIController.py @@ -1,8 +1,6 @@ import os -import platform import subprocess import traceback -from tkinter import filedialog from modules.module.Blip2Model import Blip2Model from modules.module.BlipModel import BlipModel @@ -11,39 +9,22 @@ from modules.module.RembgHumanModel import RembgHumanModel from modules.module.RembgModel import RembgModel from modules.module.WDModel import WDModel -from modules.ui.GenerateCaptionsWindow import GenerateCaptionsWindow -from modules.ui.GenerateMasksWindow import GenerateMasksWindow +from modules.ui.GenerateCaptionsWindowController import GenerateCaptionsWindowController +from modules.ui.GenerateMasksWindowController import GenerateMasksWindowController from modules.util import path_util from modules.util.image_util import load_image from modules.util.torch_util import default_device, torch_gc -from modules.util.ui import components -from modules.util.ui.ui_utils import bind_mousewheel, set_window_icon -from modules.util.ui.UIState import UIState import torch -import customtkinter as ctk -import cv2 import numpy as np -from customtkinter import ScalingTracker, ThemeManager from PIL import Image, ImageDraw -class CaptionUI(ctk.CTkToplevel): - def __init__( - self, - parent, - initial_dir: str | None, - initial_include_subdirectories: bool, - *args, - **kwargs, - ) -> None: - super().__init__(parent, *args, **kwargs) - self.protocol("WM_DELETE_WINDOW", self._on_close) - +class CaptionUIController: + def __init__(self, initial_dir: str | None, initial_include_subdirectories: bool): self.dir = initial_dir self.config_ui_data = {"include_subdirectories": initial_include_subdirectories} - self.config_ui_state = UIState(self, self.config_ui_data) self.image_size = 850 self.help_text = """ Keyboard shortcuts when focusing on the prompt input field: @@ -62,8 +43,6 @@ def __init__( self.captioning_model = None self.image_rel_paths = [] self.current_image_index = -1 - self.file_list = None - self.image_labels = [] self.pil_image = None self.image_width = 0 self.image_height = 0 @@ -72,155 +51,67 @@ def __init__( self.mask_draw_y = 0 self.mask_draw_radius = 0.01 self.display_only_mask = False - self.image = None - self.image_label = None self.mask_editing_mode = 'draw' - self.enable_mask_editing_var = ctk.BooleanVar() - self.mask_editing_alpha = None - self.prompt_var = None - self.prompt_component = None - + self.view = None - self.title("OneTrainer") - self.geometry("1280x980") - self.resizable(False, False) + def create_window(self, parent, view_cls): + self.view = view_cls(parent, self) + return self.view + def open_mask_window(self, parent_window, view_cls): + controller = GenerateMasksWindowController(self) + return controller.create_window(parent_window, self.dir, self.config_ui_data["include_subdirectories"], view_cls) - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=1) - self.grid_columnconfigure(0, weight=1) + def open_caption_window(self, parent_window, view_cls): + controller = GenerateCaptionsWindowController(self) + return controller.create_window(parent_window, self.dir, self.config_ui_data["include_subdirectories"], view_cls) + def open_in_explorer(self): + try: + image_name = self.image_rel_paths[self.current_image_index] + image_name = os.path.realpath(os.path.join(self.dir, image_name)) + subprocess.Popen(f"explorer /select,{image_name}") + except Exception: + traceback.print_exc() - self.top_bar(self) + def switch_image(self, index): + old_index = self.current_image_index + self.current_image_index = index + if index >= 0: + self.pil_image = self.load_image() + self.pil_mask = self.load_mask() + prompt = self.load_prompt() - self.bottom_frame = ctk.CTkFrame(self) - self.bottom_frame.grid(row=1, column=0, sticky="nsew") - self.bottom_frame.grid_rowconfigure(0, weight=1) - self.bottom_frame.grid_columnconfigure(0, weight=0) - self.bottom_frame.grid_columnconfigure(1, weight=1) + self.image_width = self.pil_image.width + self.image_height = self.pil_image.height + scale = self.image_size / max(self.pil_image.height, self.pil_image.width) + height = int(self.pil_image.height * scale) + width = int(self.pil_image.width * scale) - self.file_list_column(self.bottom_frame) - self.content_column(self.bottom_frame) - self.load_directory() + self.pil_image = self.pil_image.resize((width, height), Image.Resampling.LANCZOS) - self.wait_visibility() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - def top_bar(self, master): - top_frame = ctk.CTkFrame(master) - top_frame.grid(row=0, column=0, sticky="nsew") - - components.button(top_frame, 0, 0, "Open", self.open_directory, - tooltip="open a new directory") - components.button(top_frame, 0, 1, "Generate Masks", self.open_mask_window, - tooltip="open a dialog to automatically generate masks") - components.button(top_frame, 0, 2, "Generate Captions", self.open_caption_window, - tooltip="open a dialog to automatically generate captions") - - if platform.system() == "Windows": - components.button(top_frame, 0, 3, "Open in Explorer", self.open_in_explorer, - tooltip="open the current image in Explorer") - - components.switch(top_frame, 0, 4, self.config_ui_state, "include_subdirectories", - text="include subdirectories") - - top_frame.grid_columnconfigure(5, weight=1) - - components.button(top_frame, 0, 6, "Help", self.print_help, - tooltip=self.help_text) - - def file_list_column(self, master): - if self.file_list is not None: - self.image_labels = [] - self.file_list.destroy() - - self.file_list = ctk.CTkScrollableFrame(master, width=300) - self.file_list.grid(row=0, column=0, sticky="nsew") - - for i, filename in enumerate(self.image_rel_paths): - def __create_switch_image(index): - def __switch_image(event): - self.switch_image(index) - - return __switch_image - - label = ctk.CTkLabel(self.file_list, text=filename) - label.bind("", __create_switch_image(i)) - - self.image_labels.append(label) - label.grid(row=i, column=0, padx=5, sticky="nsw") - - def content_column(self, master): - image = Image.new("RGBA", (512, 512), (0, 0, 0, 0)) - - right_frame = ctk.CTkFrame(master, fg_color="transparent") - right_frame.grid(row=0, column=1, sticky="nsew") - - right_frame.grid_columnconfigure(4, weight=1) - right_frame.grid_rowconfigure(1, weight=1) - - components.button(right_frame, 0, 0, "Draw", self.draw_mask_editing_mode, - tooltip="draw a mask using a brush") - components.button(right_frame, 0, 1, "Fill", self.fill_mask_editing_mode, - tooltip="draw a mask using a fill tool") - - # checkbox to enable mask editing - self.enable_mask_editing_var = ctk.BooleanVar() - self.enable_mask_editing_var.set(False) - enable_mask_editing_checkbox = ctk.CTkCheckBox( - right_frame, text="Enable Mask Editing", variable=self.enable_mask_editing_var, width=50) - enable_mask_editing_checkbox.grid(row=0, column=2, padx=25, pady=5, sticky="w") - - # mask alpha textbox - self.mask_editing_alpha = ctk.CTkEntry(master=right_frame, width=40, placeholder_text="1.0") - self.mask_editing_alpha.insert(0, "1.0") - self.mask_editing_alpha.grid(row=0, column=3, sticky="e", padx=5, pady=5) - self.bind_key_events(self.mask_editing_alpha) + self.view.on_image_switched(old_index, index, prompt) + else: + self.view.on_image_cleared() - mask_editing_alpha_label = ctk.CTkLabel(right_frame, text="Brush Alpha", width=75) - mask_editing_alpha_label.grid(row=0, column=4, padx=0, pady=5, sticky="w") + def previous_image(self): + if len(self.image_rel_paths) > 0 and (self.current_image_index - 1) >= 0: + self.view.switch_image(self.current_image_index - 1) - # image - self.image = ctk.CTkImage( - light_image=image, - size=(self.image_size, self.image_size) - ) - self.image_label = ctk.CTkLabel( - master=right_frame, text="", image=self.image, height=self.image_size, width=self.image_size - ) - self.image_label.grid(row=1, column=0, columnspan=5, sticky="nsew") - - self.image_label.bind("", self.edit_mask) - self.image_label.bind("", self.edit_mask) - self.image_label.bind("", self.edit_mask) - bind_mousewheel(self.image_label, {self.image_label.children["!label"]}, self.draw_mask_radius) - - # prompt - self.prompt_var = ctk.StringVar() - self.prompt_component = ctk.CTkEntry(right_frame, textvariable=self.prompt_var) - self.prompt_component.grid(row=2, column=0, columnspan=5, pady=5, sticky="new") - self.bind_key_events(self.prompt_component) - self.prompt_component.focus_set() - - def bind_key_events(self, component): - component.bind("", self.next_image) - component.bind("", self.previous_image) - component.bind("", self.save) - component.bind("", self.toggle_mask) - component.bind("", self.draw_mask_editing_mode) - component.bind("", self.fill_mask_editing_mode) + def next_image(self): + if len(self.image_rel_paths) > 0 and (self.current_image_index + 1) < len(self.image_rel_paths): + self.view.switch_image(self.current_image_index + 1) def load_directory(self, include_subdirectories: bool = False): self.scan_directory(include_subdirectories) - self.file_list_column(self.bottom_frame) + self.view.refresh_file_list() if len(self.image_rel_paths) > 0: self.switch_image(0) else: self.switch_image(-1) - self.prompt_component.focus_set() + self.view.focus_prompt() def scan_directory(self, include_subdirectories: bool = False): def __is_supported_image_extension(filename): @@ -285,42 +176,26 @@ def load_prompt(self): else: return "" - def previous_image(self, event): - if len(self.image_rel_paths) > 0 and (self.current_image_index - 1) >= 0: - self.switch_image(self.current_image_index - 1) - - def next_image(self, event): - if len(self.image_rel_paths) > 0 and (self.current_image_index + 1) < len(self.image_rel_paths): - self.switch_image(self.current_image_index + 1) - - def switch_image(self, index): - if len(self.image_labels) > 0 and self.current_image_index < len(self.image_labels): - self.image_labels[self.current_image_index].configure( - text_color=ThemeManager.theme["CTkLabel"]["text_color"]) - - self.current_image_index = index - if index >= 0: - self.image_labels[index].configure(text_color="#FF0000") + def save(self, prompt_text): + if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): + image_name = self.image_rel_paths[self.current_image_index] - self.pil_image = self.load_image() - self.pil_mask = self.load_mask() - prompt = self.load_prompt() + prompt_name = os.path.splitext(image_name)[0] + ".txt" + prompt_name = os.path.join(self.dir, prompt_name) - self.image_width = self.pil_image.width - self.image_height = self.pil_image.height - scale = self.image_size / max(self.pil_image.height, self.pil_image.width) - height = int(self.pil_image.height * scale) - width = int(self.pil_image.width * scale) + mask_name = os.path.splitext(image_name)[0] + "-masklabel.png" + mask_name = os.path.join(self.dir, mask_name) - self.pil_image = self.pil_image.resize((width, height), Image.Resampling.LANCZOS) + try: + with open(prompt_name, "w", encoding='utf-8') as f: + f.write(prompt_text) + except Exception: + return - self.refresh_image() - self.prompt_var.set(prompt) - else: - image = Image.new("RGB", (512, 512), (0, 0, 0)) - self.image.configure(light_image=image) + if self.pil_mask: + self.pil_mask.save(mask_name) - def refresh_image(self): + def get_display_image(self): if self.pil_mask: resized_pil_mask = self.pil_mask.resize( (self.pil_image.width, self.pil_image.height), @@ -328,7 +203,7 @@ def refresh_image(self): ) if self.display_only_mask: - self.image.configure(light_image=resized_pil_mask, size=resized_pil_mask.size) + return resized_pil_mask, resized_pil_mask.size else: np_image = np.array(self.pil_image).astype(np.float32) / 255.0 np_mask = np.array(resized_pil_mask).astype(np.float32) / 255.0 @@ -346,29 +221,26 @@ def refresh_image(self): np_masked_image = (np_image * np_mask * 255.0).astype(np.uint8) masked_image = Image.fromarray(np_masked_image, mode='RGB') - self.image.configure(light_image=masked_image, size=masked_image.size) + return masked_image, masked_image.size else: - self.image.configure(light_image=self.pil_image, size=self.pil_image.size) + return self.pil_image, self.pil_image.size + + def toggle_mask(self): + self.display_only_mask = not self.display_only_mask + + def set_mask_editing_mode(self, mode): + self.mask_editing_mode = mode - def draw_mask_radius(self, delta, raw_event): + def update_mask_draw_radius(self, delta): # Wheel up = Increase radius. Wheel down = Decrease radius. multiplier = 1.0 + (delta * 0.05) self.mask_draw_radius = max(0.0025, self.mask_draw_radius * multiplier) - def edit_mask(self, event): - if not self.enable_mask_editing_var.get(): - return - - if event.widget != self.image_label.children["!label"]: - return - + def handle_edit_mask(self, event_x, event_y, is_left, is_right, alpha): if len(self.image_rel_paths) == 0 or self.current_image_index >= len(self.image_rel_paths): return - - display_scaling = ScalingTracker.get_window_scaling(self) - - event_x = event.x / display_scaling - event_y = event.y / display_scaling + if self.pil_image is None: + return start_x = int(event_x / self.pil_image.width * self.image_width) start_y = int(event_y / self.pil_image.height * self.image_height) @@ -378,27 +250,16 @@ def edit_mask(self, event): self.mask_draw_x = event_x self.mask_draw_y = event_y - is_right = False - is_left = False - if event.state & 0x0100 or event.num == 1: # left mouse button - is_left = True - elif event.state & 0x0400 or event.num == 3: # right mouse button - is_right = True - if self.mask_editing_mode == 'draw': - self.draw_mask(start_x, start_y, end_x, end_y, is_left, is_right) + self.draw_mask(start_x, start_y, end_x, end_y, is_left, is_right, alpha) if self.mask_editing_mode == 'fill': - self.fill_mask(start_x, start_y, end_x, end_y, is_left, is_right) + self.fill_mask(start_x, start_y, end_x, end_y, is_left, is_right, alpha) - def draw_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): + def draw_mask(self, start_x, start_y, end_x, end_y, is_left, is_right, alpha): color = None adding_to_mask = True if is_left: - try: - alpha = float(self.mask_editing_alpha.get()) - except Exception: - alpha = 1.0 rgb_value = int(max(0.0, min(alpha, 1.0)) * 255) # max/min stuff to clamp to 0 - 255 range color = (rgb_value, rgb_value, rgb_value) @@ -423,17 +284,13 @@ def draw_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): draw.ellipse((end_x - radius, end_y - radius, end_x + radius, end_y + radius), fill=color, outline=None) - self.refresh_image() + self.view.refresh_image() - def fill_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): + def fill_mask(self, start_x, start_y, end_x, end_y, is_left, is_right, alpha): color = None adding_to_mask = True if is_left: - try: - alpha = float(self.mask_editing_alpha.get()) - except Exception: - alpha = 1.0 rgb_value = int(max(0.0, min(alpha, 1.0)) * 255) # max/min stuff to clamp to 0 - 255 range color = (rgb_value, rgb_value, rgb_value) @@ -449,69 +306,11 @@ def fill_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(255, 255, 255)) np_mask = np.array(self.pil_mask).astype(np.uint8) + import cv2 cv2.floodFill(np_mask, None, (start_x, start_y), color) self.pil_mask = Image.fromarray(np_mask, 'RGB') - self.refresh_image() - - def save(self, event): - if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): - image_name = self.image_rel_paths[self.current_image_index] - - prompt_name = os.path.splitext(image_name)[0] + ".txt" - prompt_name = os.path.join(self.dir, prompt_name) - - mask_name = os.path.splitext(image_name)[0] + "-masklabel.png" - mask_name = os.path.join(self.dir, mask_name) - - try: - with open(prompt_name, "w", encoding='utf-8') as f: - f.write(self.prompt_var.get()) - except Exception: - return - - if self.pil_mask: - self.pil_mask.save(mask_name) - - def draw_mask_editing_mode(self, *args): - self.mask_editing_mode = 'draw' - - if args: - # disable default event - return "break" - return None - - def fill_mask_editing_mode(self, *args): - self.mask_editing_mode = 'fill' - - def toggle_mask(self, *args): - self.display_only_mask = not self.display_only_mask - self.refresh_image() - - def open_directory(self): - new_dir = filedialog.askdirectory() - - if new_dir: - self.dir = new_dir - self.load_directory(include_subdirectories=self.config_ui_data["include_subdirectories"]) - - def open_mask_window(self): - dialog = GenerateMasksWindow(self, self.dir, self.config_ui_data["include_subdirectories"]) - self.wait_window(dialog) - self.switch_image(self.current_image_index) - - def open_caption_window(self): - dialog = GenerateCaptionsWindow(self, self.dir, self.config_ui_data["include_subdirectories"]) - self.wait_window(dialog) - self.switch_image(self.current_image_index) - - def open_in_explorer(self): - try: - image_name = self.image_rel_paths[self.current_image_index] - image_name = os.path.realpath(os.path.join(self.dir, image_name)) - subprocess.Popen(f"explorer /select,{image_name}") - except Exception: - traceback.print_exc() + self.view.refresh_image() def load_masking_model(self, model): model_type = type(self.masking_model).__name__ if self.masking_model else None @@ -563,10 +362,6 @@ def _release_models(self): if freed: torch_gc() - def _on_close(self): - self._release_models() - self.destroy() - - def destroy(self): + def on_close(self): self._release_models() - super().destroy() + self.view.destroy() diff --git a/modules/ui/CloudTabController.py b/modules/ui/CloudTabController.py new file mode 100644 index 000000000..d21dda6ce --- /dev/null +++ b/modules/ui/CloudTabController.py @@ -0,0 +1,31 @@ + +import webbrowser + +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.CloudType import CloudType + + +class CloudTabController: + def __init__(self, config: TrainConfig, parent): + self.config = config + self.parent = parent + self.reattach = False + + def do_reattach(self): + self.reattach = True + try: + self.parent.start_training() + finally: + self.reattach = False + + def get_gpu_types(self) -> list[str]: + if self.config.cloud.type == CloudType.RUNPOD: + import runpod + runpod.api_key = self.config.secrets.cloud.api_key + gpus = runpod.get_gpus() + return [gpu['id'] for gpu in gpus] + return [] + + def open_create_cloud_url(self): + if self.config.cloud.type == CloudType.RUNPOD: + webbrowser.open("https://www.runpod.io/console/deploy?template=1a33vbssq9&type=gpu", new=0, autoraise=False) diff --git a/modules/ui/ConceptTabController.py b/modules/ui/ConceptTabController.py new file mode 100644 index 000000000..b90b7a4f6 --- /dev/null +++ b/modules/ui/ConceptTabController.py @@ -0,0 +1,19 @@ + +from modules.ui.ConceptWindowController import ConceptWindowController +from modules.util.config.ConceptConfig import ConceptConfig +from modules.util.config.TrainConfig import TrainConfig + + +class ConceptTabController: + def __init__(self, config: TrainConfig): + self.train_config = config + + def create_new_element(self) -> ConceptConfig: + return ConceptConfig.default_values() + + def randomize_seed(self, concept: ConceptConfig) -> ConceptConfig: + concept.seed = ConceptConfig.default_values().seed + return concept + + def open_element_window(self, parent, concept_config, ui_state, image_ui_state, text_ui_state, view_cls): + return view_cls(parent, ConceptWindowController(self.train_config, concept_config), ui_state, image_ui_state, text_ui_state) diff --git a/modules/ui/ConceptWindowController.py b/modules/ui/ConceptWindowController.py index f58879d5f..c2c486961 100644 --- a/modules/ui/ConceptWindowController.py +++ b/modules/ui/ConceptWindowController.py @@ -1,5 +1,3 @@ -import fractions -import math import os import pathlib import platform @@ -11,12 +9,7 @@ from modules.util import concept_stats, path_util from modules.util.config.ConceptConfig import ConceptConfig from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.BalancingStrategy import BalancingStrategy -from modules.util.enum.ConceptType import ConceptType from modules.util.image_util import load_image -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState from mgds.LoadingPipeline import LoadingPipeline from mgds.OutputPipelineModule import OutputPipelineModule @@ -34,534 +27,21 @@ from mgds.pipelineModules.RandomRotate import RandomRotate from mgds.pipelineModules.RandomSaturation import RandomSaturation from mgds.pipelineModules.ShuffleTags import ShuffleTags -from mgds.pipelineModuleTypes.RandomAccessPipelineModule import ( - RandomAccessPipelineModule, -) +from mgds.pipelineModuleTypes.RandomAccessPipelineModule import RandomAccessPipelineModule import torch from torchvision.transforms import functional -import customtkinter as ctk 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 -class InputPipelineModule( - PipelineModule, - RandomAccessPipelineModule, -): - def __init__(self, data: dict): - super().__init__() - self.data = data - - def length(self) -> int: - return 1 - - def get_inputs(self) -> list[str]: - return [] - - def get_outputs(self) -> list[str]: - return list(self.data.keys()) - - def get_item(self, variation: int, index: int, requested_name: str = None) -> dict: - return self.data - - -class ConceptWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - train_config: TrainConfig, - concept: ConceptConfig, - ui_state: UIState, - image_ui_state: UIState, - text_ui_state: UIState, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) - +class ConceptWindowController: + def __init__(self, train_config: TrainConfig, concept: ConceptConfig): self.train_config = train_config - self.concept = concept - self.ui_state = ui_state - self.image_ui_state = image_ui_state - 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.title("Concept") - self.geometry("800x700") - self.resizable(True, True) - - self.grid_rowconfigure(0, weight=1) - self.grid_columnconfigure(0, weight=1) - - tabview = ctk.CTkTabview(self) - tabview.grid(row=0, column=0, sticky="nsew") - - self.general_tab = self.__general_tab(tabview.add("general"), concept) - self.image_augmentation_tab = self.__image_augmentation_tab(tabview.add("image augmentation")) - self.text_augmentation_tab = self.__text_augmentation_tab(tabview.add("text augmentation")) - self.concept_stats_tab = self.__concept_stats_tab(tabview.add("statistics")) - - #automatic concept scan - self.scan_thread = threading.Thread(target=self.__auto_update_concept_stats, daemon=True) - self.scan_thread.start() - - components.button(self, 1, 0, "ok", self.__ok) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def __general_tab(self, master, concept: ConceptConfig): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, weight=1) - - # name - components.label(frame, 0, 0, "Name", - tooltip="Name of the concept") - components.entry(frame, 0, 1, self.ui_state, "name") - - # enabled - components.label(frame, 1, 0, "Enabled", - tooltip="Enable or disable this concept") - components.switch(frame, 1, 1, self.ui_state, "enabled") - - # concept type - components.label(frame, 2, 0, "Concept Type", - tooltip="STANDARD: Standard finetuning with the sample as training target\n" - "VALIDATION: Use concept for validation instead of training\n" - "PRIOR_PREDICTION: Use the sample to make a prediction using the model as it was before training. This prediction is then used as the training target " - "for the model in training. This can be used as regularisation and to preserve prior model knowledge while finetuning the model on other concepts. " - "Only implemented for LoRA.", - wide_tooltip=True) - components.options(frame, 2, 1, [str(x) for x in list(ConceptType)], self.ui_state, "type") - - # path - components.label(frame, 3, 0, "Path", - tooltip="Path where the training data is located") - components.path_entry(frame, 3, 1, self.ui_state, "path", mode="dir") - 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.path_entry(frame, 4, 2, self.text_ui_state, "prompt_path", mode="file") - - def set_prompt_path_entry_enabled(option: str): - if option == 'concept': - for child in prompt_path_entry.children.values(): - child.configure(state="normal") - else: - for child in prompt_path_entry.children.values(): - child.configure(state="disabled") - - components.options_kv(frame, 4, 1, [ - ("From text file per sample", 'sample'), - ("From single text file", 'concept'), - ("From image file name", 'filename'), - ], self.text_ui_state, "prompt_source", command=set_prompt_path_entry_enabled) - set_prompt_path_entry_enabled(concept.text.prompt_source) - - # include subdirectories - components.label(frame, 5, 0, "Include Subdirectories", - tooltip="Includes images from subdirectories into the dataset") - components.switch(frame, 5, 1, self.ui_state, "include_subdirectories") - - # image variations - components.label(frame, 6, 0, "Image Variations", - tooltip="The number of different image versions to cache if latent caching is enabled.") - components.entry(frame, 6, 1, self.ui_state, "image_variations") - - # text variations - components.label(frame, 7, 0, "Text Variations", - tooltip="The number of different text versions to cache if latent caching is enabled.") - components.entry(frame, 7, 1, self.ui_state, "text_variations") - - # balancing - components.label(frame, 8, 0, "Balancing", - tooltip="The number of samples used during training. Use repeats to multiply the concept, or samples to specify an exact number of samples used in each epoch.") - components.entry(frame, 8, 1, self.ui_state, "balancing") - components.options(frame, 8, 2, [str(x) for x in list(BalancingStrategy)], self.ui_state, "balancing_strategy") - - # loss weight - components.label(frame, 9, 0, "Loss Weight", - tooltip="The loss multiplyer for this concept.") - components.entry(frame, 9, 1, self.ui_state, "loss_weight") - - frame.pack(fill="both", expand=1) - return frame - - def __image_augmentation_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # header - components.label(frame, 0, 1, "Random", - tooltip="Enable this augmentation with random values") - components.label(frame, 0, 2, "Fixed", - tooltip="Enable this augmentation with fixed values") - - # crop jitter - components.label(frame, 1, 0, "Crop Jitter", - tooltip="Enables random cropping of samples") - components.switch(frame, 1, 1, self.image_ui_state, "enable_crop_jitter") - - # random flip - components.label(frame, 2, 0, "Random Flip", - tooltip="Randomly flip the sample during training") - components.switch(frame, 2, 1, self.image_ui_state, "enable_random_flip") - components.switch(frame, 2, 2, self.image_ui_state, "enable_fixed_flip") - - # random rotation - components.label(frame, 3, 0, "Random Rotation", - tooltip="Randomly rotates the sample during training") - components.switch(frame, 3, 1, self.image_ui_state, "enable_random_rotate") - components.switch(frame, 3, 2, self.image_ui_state, "enable_fixed_rotate") - components.entry(frame, 3, 3, self.image_ui_state, "random_rotate_max_angle") - - # random brightness - components.label(frame, 4, 0, "Random Brightness", - tooltip="Randomly adjusts the brightness of the sample during training") - components.switch(frame, 4, 1, self.image_ui_state, "enable_random_brightness") - components.switch(frame, 4, 2, self.image_ui_state, "enable_fixed_brightness") - components.entry(frame, 4, 3, self.image_ui_state, "random_brightness_max_strength") - - # random contrast - components.label(frame, 5, 0, "Random Contrast", - tooltip="Randomly adjusts the contrast of the sample during training") - components.switch(frame, 5, 1, self.image_ui_state, "enable_random_contrast") - components.switch(frame, 5, 2, self.image_ui_state, "enable_fixed_contrast") - components.entry(frame, 5, 3, self.image_ui_state, "random_contrast_max_strength") - - # random saturation - components.label(frame, 6, 0, "Random Saturation", - tooltip="Randomly adjusts the saturation of the sample during training") - components.switch(frame, 6, 1, self.image_ui_state, "enable_random_saturation") - components.switch(frame, 6, 2, self.image_ui_state, "enable_fixed_saturation") - components.entry(frame, 6, 3, self.image_ui_state, "random_saturation_max_strength") - - # random hue - components.label(frame, 7, 0, "Random Hue", - tooltip="Randomly adjusts the hue of the sample during training") - components.switch(frame, 7, 1, self.image_ui_state, "enable_random_hue") - components.switch(frame, 7, 2, self.image_ui_state, "enable_fixed_hue") - components.entry(frame, 7, 3, self.image_ui_state, "random_hue_max_strength") - - # random circular mask shrink - components.label(frame, 8, 0, "Circular Mask Generation", - tooltip="Automatically create circular masks for masked training") - components.switch(frame, 8, 1, self.image_ui_state, "enable_random_circular_mask_shrink") - - # random rotate and crop - components.label(frame, 9, 0, "Random Rotate and Crop", - tooltip="Randomly rotate the training samples and crop to the masked region") - components.switch(frame, 9, 1, self.image_ui_state, "enable_random_mask_rotate_crop") - - # circular mask generation - components.label(frame, 10, 0, "Resolution Override", - tooltip="Override the resolution for this concept. Optionally specify multiple resolutions separated by a comma, or a single exact resolution in the format x") - components.switch(frame, 10, 2, self.image_ui_state, "enable_resolution_override") - components.entry(frame, 10, 3, self.image_ui_state, "resolution_override") - - # image - image_preview, filename_preview, caption_preview = self.__get_preview_image() - self.image = ctk.CTkImage( - light_image=image_preview, - size=image_preview.size, - ) - image_label = ctk.CTkLabel(master=frame, text="", image=self.image, height=300, width=300) - image_label.grid(row=0, column=4, rowspan=6) - - # refresh preview - update_button_frame = ctk.CTkFrame(master=frame, corner_radius=0, fg_color="transparent") - update_button_frame.grid(row=6, column=4, rowspan=6, sticky="nsew") - update_button_frame.grid_columnconfigure(1, weight=1) - - prev_preview_button = components.button(update_button_frame, 0, 0, "<", command=self.__prev_image_preview) - components.button(update_button_frame, 0, 1, "Update Preview", command=self.__update_image_preview) - next_preview_button = components.button(update_button_frame, 0, 2, ">", command=self.__next_image_preview) - preview_augmentations_switch = ctk.CTkSwitch(update_button_frame, text="Show Augmentations", variable=self.preview_augmentations, command=self.__update_image_preview) - preview_augmentations_switch.grid(row=1, column=0, columnspan=3, padx=5, pady=5) - - prev_preview_button.configure(width=40) - next_preview_button.configure(width=40) - - #caption and filename preview - self.filename_preview = ctk.CTkLabel(master=update_button_frame, text=filename_preview, width=300, anchor="nw", justify="left", padx=10, wraplength=280) - self.filename_preview.grid(row=2, column=0, columnspan=3) - self.caption_preview = ctk.CTkTextbox(master=update_button_frame, width = 300, height = 150, wrap="word", border_width=2) - self.caption_preview.insert(index="1.0", text=caption_preview) - self.caption_preview.configure(state="disabled") - self.caption_preview.grid(row=3, column=0, columnspan=3, rowspan=3) - - frame.pack(fill="both", expand=1) - return frame - - def __text_augmentation_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # tag shuffling - components.label(frame, 0, 0, "Tag Shuffling", - tooltip="Enables tag shuffling") - components.switch(frame, 0, 1, self.text_ui_state, "enable_tag_shuffling") - - # keep tag count - components.label(frame, 1, 0, "Tag Delimiter", - tooltip="The delimiter between tags") - components.entry(frame, 1, 1, self.text_ui_state, "tag_delimiter") - - # keep tag count - components.label(frame, 2, 0, "Keep Tag Count", - tooltip="The number of tags at the start of the caption that are not shuffled or dropped") - components.entry(frame, 2, 1, self.text_ui_state, "keep_tags_count") - - # tag dropout - components.label(frame, 3, 0, "Tag Dropout", - tooltip="Enables random dropout for tags in the captions.") - components.switch(frame, 3, 1, self.text_ui_state, "tag_dropout_enable") - components.label(frame, 4, 0, "Dropout Mode", - tooltip="Method used to drop captions. 'Full' will drop the entire caption past the 'kept' tags with a certain probability, 'Random' will drop individual tags with the set probability, and 'Random Weighted' will linearly increase the probability of dropping tags, more likely to preseve tags near the front with full probability to drop at the end.") - components.options_kv(frame, 4, 1, [ - ("Full", 'FULL'), - ("Random", 'RANDOM'), - ("Random Weighted", 'RANDOM WEIGHTED'), - ], self.text_ui_state, "tag_dropout_mode", None) - components.label(frame, 4, 2, "Probability", - tooltip="Probability to drop tags, from 0 to 1.") - components.entry(frame, 4, 3, self.text_ui_state, "tag_dropout_probability") - - components.label(frame, 5, 0, "Special Dropout Tags", - tooltip="List of tags which will be whitelisted/blacklisted by dropout. 'Whitelist' tags will never be dropped but all others may be, 'Blacklist' tags may be dropped but all others will never be, 'None' may drop any tags. Can specify either a delimiter-separated list in the field, or a file path to a .txt or .csv file with entries separated by newlines.") - components.options_kv(frame, 5, 1, [ - ("None", 'NONE'), - ("Blacklist", 'BLACKLIST'), - ("Whitelist", 'WHITELIST'), - ], self.text_ui_state, "tag_dropout_special_tags_mode", None) - components.entry(frame, 5, 2, self.text_ui_state, "tag_dropout_special_tags") - components.label(frame, 6, 0, "Special Tags Regex", - tooltip="Interpret special tags with regex, such as 'photo.*' to match 'photo, photograph, photon' but not 'telephoto'. Includes exception for '/(' and '/)' syntax found in many booru/e6 tags.") - components.switch(frame, 6, 1, self.text_ui_state, "tag_dropout_special_tags_regex") - - #capitalization randomization - components.label(frame, 7, 0, "Randomize Capitalization", - tooltip="Enables randomization of capitalization for tags in the caption.") - components.switch(frame, 7, 1, self.text_ui_state, "caps_randomize_enable") - components.label(frame, 7, 2, "Force Lowercase", - tooltip="If enabled, converts the caption to lowercase before any further processing.") - components.switch(frame, 7, 3, self.text_ui_state, "caps_randomize_lowercase") - - components.label(frame, 8, 0, "Captialization Mode", - tooltip="Comma-separated list of types of capitalization randomization to perform. 'capslock' for ALL CAPS, 'title' for First Letter Of Every Word, 'first' for First word only, 'random' for rAndOMiZeD lEtTERs.") - components.entry(frame, 8, 1, self.text_ui_state, "caps_randomize_mode") - components.label(frame, 8, 2, "Probability", - tooltip="Probability to randomize capitialization of each tag, from 0 to 1.") - components.entry(frame, 8, 3, self.text_ui_state, "caps_randomize_probability") - - frame.pack(fill="both", expand=1) - return frame - - def __concept_stats_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=150) - frame.grid_columnconfigure(1, weight=0, minsize=150) - frame.grid_columnconfigure(2, weight=0, minsize=150) - frame.grid_columnconfigure(3, weight=0, minsize=150) - self.cancel_scan_flag = threading.Event() - - #file size - self.file_size_label = components.label(frame, 1, 0, "Total Size", pad=0, - tooltip="Total size of all image, mask, and caption files in MB") - self.file_size_label.configure(font=ctk.CTkFont(underline=True)) - self.file_size_preview = components.label(frame, 2, 0, pad=0, text="-") - - #subdirectory count - self.dir_count_label = components.label(frame, 1, 1, "Directories", pad=0, - tooltip="Total number of directories including and under (if 'include subdirectories' is enabled) the main concept directory") - self.dir_count_label.configure(font=ctk.CTkFont(underline=True)) - self.dir_count_preview = components.label(frame, 2, 1, pad=0, text="-") - - #basic img/vid stats - count of each type in the concept - #the \n at the start of the label gives it better vertical spacing with other rows - self.image_count_label = components.label(frame, 3, 0, "\nTotal Images", pad=0, - tooltip="Total number of image files, any of the extensions " + str(path_util.SUPPORTED_IMAGE_EXTENSIONS) + ", excluding '-masklabel.png and -condlabel.png'") - self.image_count_label.configure(font=ctk.CTkFont(underline=True)) - self.image_count_preview = components.label(frame, 4, 0, pad=0, text="-") - self.video_count_label = components.label(frame, 3, 1, "\nTotal Videos", pad=0, - tooltip="Total number of video files, any of the extensions " + str(path_util.SUPPORTED_VIDEO_EXTENSIONS)) - self.video_count_label.configure(font=ctk.CTkFont(underline=True)) - self.video_count_preview = components.label(frame, 4, 1, pad=0, text="-") - self.mask_count_label = components.label(frame, 3, 2, "\nTotal Masks", pad=0, - tooltip="Total number of mask files, any file ending in '-masklabel.png'") - self.mask_count_label.configure(font=ctk.CTkFont(underline=True)) - self.mask_count_preview = components.label(frame, 4, 2, pad=0, text="-") - self.caption_count_label = components.label(frame, 3, 3, "\nTotal Captions", pad=0, - tooltip="Total number of caption files, any .txt file. With advanced scan, includes the total number of captions on separate lines across all files in parentheses.") - self.caption_count_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_count_preview = components.label(frame, 4, 3, pad=0, text="-") - - #advanced img/vid stats - how many img/vid files have a mask or caption of the same name - self.image_count_mask_label = components.label(frame, 5, 0, "\nImages with Masks", pad=0, - tooltip="Total number of image files with an associated mask") - self.image_count_mask_label.configure(font=ctk.CTkFont(underline=True)) - self.image_count_mask_preview = components.label(frame, 6, 0, pad=0, text="-") - self.mask_count_label_unpaired = components.label(frame, 5, 1, "\nUnpaired Masks", pad=0, - tooltip="Total number of mask files which lack a corresponding image file - if >0, check your data set!") - self.mask_count_label_unpaired.configure(font=ctk.CTkFont(underline=True)) - self.mask_count_preview_unpaired = components.label(frame, 6, 1, pad=0, text="-") - #currently no masks for videos? - - self.image_count_caption_label = components.label(frame, 7, 0, "\nImages with Captions", pad=0, - tooltip="Total number of image files with an associated caption") - self.image_count_caption_label.configure(font=ctk.CTkFont(underline=True)) - self.image_count_caption_preview = components.label(frame, 8, 0, pad=0, text="-") - self.video_count_caption_label = components.label(frame, 7, 1, "\nVideos with Captions", pad=0, - tooltip="Total number of video files with an associated caption") - self.video_count_caption_label.configure(font=ctk.CTkFont(underline=True)) - self.video_count_caption_preview = components.label(frame, 8, 1, pad=0, text="-") - self.caption_count_label_unpaired = components.label(frame, 7, 2, "\nUnpaired Captions", pad=0, - tooltip="Total number of caption files which lack a corresponding image file - if >0, check your data set! If using 'from file name' or 'from single text file' then this can be ignored.") - self.caption_count_label_unpaired.configure(font=ctk.CTkFont(underline=True)) - self.caption_count_preview_unpaired = components.label(frame, 8, 2, pad=0, text="-") - - #resolution info - self.pixel_max_label = components.label(frame, 9, 0, "\nMax Pixels", pad=0, - tooltip="Largest image in the concept by number of pixels (width * height)") - self.pixel_max_label.configure(font=ctk.CTkFont(underline=True)) - self.pixel_max_preview = components.label(frame, 10, 0, pad=0, text="-", wraplength=150) - self.pixel_avg_label = components.label(frame, 9, 1, "\nAvg Pixels", pad=0, - tooltip="Average size of images in the concept by number of pixels (width * height)") - self.pixel_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.pixel_avg_preview = components.label(frame, 10, 1, pad=0, text="-", wraplength=150) - self.pixel_min_label = components.label(frame, 9, 2, "\nMin Pixels", pad=0, - tooltip="Smallest image in the concept by number of pixels (width * height)") - self.pixel_min_label.configure(font=ctk.CTkFont(underline=True)) - self.pixel_min_preview = components.label(frame, 10, 2, pad=0, text="-", wraplength=150) - - #video length info - self.length_max_label = components.label(frame, 11, 0, "\nMax Length", pad=0, - tooltip="Longest video in the concept by number of frames") - self.length_max_label.configure(font=ctk.CTkFont(underline=True)) - self.length_max_preview = components.label(frame, 12, 0, pad=0, text="-", wraplength=150) - self.length_avg_label = components.label(frame, 11, 1, "\nAvg Length", pad=0, - tooltip="Average length of videos in the concept by number of frames") - self.length_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.length_avg_preview = components.label(frame, 12, 1, pad=0, text="-", wraplength=150) - self.length_min_label = components.label(frame, 11, 2, "\nMin Length", pad=0, - tooltip="Shortest video in the concept by number of frames") - self.length_min_label.configure(font=ctk.CTkFont(underline=True)) - self.length_min_preview = components.label(frame, 12, 2, pad=0, text="-", wraplength=150) - - #video fps info - self.fps_max_label = components.label(frame, 13, 0, "\nMax FPS", pad=0, - tooltip="Video in concept with highest fps") - self.fps_max_label.configure(font=ctk.CTkFont(underline=True)) - self.fps_max_preview = components.label(frame, 14, 0, pad=0, text="-", wraplength=150) - self.fps_avg_label = components.label(frame, 13, 1, "\nAvg FPS", pad=0, - tooltip="Average fps of videos in the concept") - self.fps_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.fps_avg_preview = components.label(frame, 14, 1, pad=0, text="-", wraplength=150) - self.fps_min_label = components.label(frame, 13, 2, "\nMin FPS", pad=0, - tooltip="Video in concept with the lowest fps") - self.fps_min_label.configure(font=ctk.CTkFont(underline=True)) - self.fps_min_preview = components.label(frame, 14, 2, pad=0, text="-", wraplength=150) - - #caption info - self.caption_max_label = components.label(frame, 15, 0, "\nMax Caption Length", pad=0, - tooltip="Largest caption in concept by character count. For token count, assume ~2 tokens/word") - self.caption_max_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_max_preview = components.label(frame, 16, 0, pad=0, text="-", wraplength=150) - self.caption_avg_label = components.label(frame, 15, 1, "\nAvg Caption Length", pad=0, - tooltip="Average length of caption in concept by character count. For token count, assume ~2 tokens/word") - self.caption_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_avg_preview = components.label(frame, 16, 1, pad=0, text="-", wraplength=150) - self.caption_min_label = components.label(frame, 15, 2, "\nMin Caption Length", pad=0, - tooltip="Smallest caption in concept by character count. For token count, assume ~2 tokens/word") - self.caption_min_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_min_preview = components.label(frame, 16, 2, pad=0, text="-", wraplength=150) - - #aspect bucket info - self.aspect_bucket_label = components.label(frame, 17, 0, "\nAspect Bucketing", pad=0, - tooltip="Graph of all possible buckets and the number of images in each one, defined as height/width. Buckets range from 0.25 (4:1 extremely wide) to 4 (1:4 extremely tall). \ - Images which don't match a bucket exactly are cropped to the nearest one.") - self.aspect_bucket_label.configure(font=ctk.CTkFont(underline=True)) - self.small_bucket_label = components.label(frame, 17, 1, "\nSmallest Buckets", pad=0, - tooltip="Image buckets with the least nonzero total images - if 'batch size' is larger than this, these images will be ignored during training! See the wiki for more details.") - self.small_bucket_label.configure(font=ctk.CTkFont(underline=True)) - self.small_bucket_preview = components.label(frame, 18, 1, pad=0, text="-") - - #aspect bucketing plot, mostly copied from timestep preview graph - appearance_mode = AppearanceModeTracker.get_mode() - background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) - text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) - background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" - self.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 - - assert self.bucket_fig is None - self.bucket_fig, self.bucket_ax = plt.subplots(figsize=(7,3)) - self.canvas = FigureCanvasTkAgg(self.bucket_fig, master=frame) - self.canvas.get_tk_widget().grid(row=19, column=0, columnspan=4, rowspan=2) - self.bucket_fig.tight_layout() - self.bucket_fig.subplots_adjust(bottom=0.15) - - self.bucket_fig.set_facecolor(background_color) - self.bucket_ax.set_facecolor(background_color) - self.bucket_ax.spines['bottom'].set_color(self.text_color) - self.bucket_ax.spines['left'].set_color(self.text_color) - self.bucket_ax.spines['top'].set_visible(False) - self.bucket_ax.spines['right'].set_color(self.text_color) - self.bucket_ax.tick_params(axis='x', colors=self.text_color, which="both") - self.bucket_ax.tick_params(axis='y', colors=self.text_color, which="both") - self.bucket_ax.xaxis.label.set_color(self.text_color) - self.bucket_ax.yaxis.label.set_color(self.text_color) - - #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") - self.refresh_advanced_stats_button = components.button(master=frame, row=0, column=1, text="Refresh Advanced", command=lambda: self.__get_concept_stats_threaded(True, 9999), - tooltip="Reload advanced statistics for the concept directory") #run "basic" scan first before "advanced", seems to help the system cache the directories and run faster - self.cancel_stats_button = components.button(master=frame, row=0, column=2, text="Abort Scan", command=lambda: self.__cancel_concept_stats(), - tooltip="Stop the currently running scan if it's taking a long time - advanced scan will be slow on large folders and on HDDs") - self.processing_time = components.label(frame, 0, 3, text="-", tooltip="Time taken to process concept directory") - - frame.pack(fill="both", expand=1) - return frame - - def __prev_image_preview(self): - self.image_preview_file_index = max(self.image_preview_file_index - 1, 0) - self.__update_image_preview() - - def __next_image_preview(self): - self.image_preview_file_index += 1 - self.__update_image_preview() - - def __update_image_preview(self): - image_preview, filename_preview, caption_preview = self.__get_preview_image() - self.image.configure(light_image=image_preview, size=image_preview.size) - self.filename_preview.configure(text=filename_preview) - self.caption_preview.configure(state="normal") - self.caption_preview.delete(index1="1.0", index2="end") - self.caption_preview.insert(index="1.0", text=caption_preview) - self.caption_preview.configure(state="disabled") + self.scan_thread = None @staticmethod def get_concept_path(path: str) -> str | None: @@ -573,22 +53,22 @@ def get_concept_path(path: str) -> str | None: except Exception: return None - def __download_dataset(self): + def download_dataset(self): try: huggingface_hub.login(token=self.train_config.secrets.huggingface_token, new_session=False) huggingface_hub.snapshot_download(repo_id=self.concept.path, repo_type="dataset") except Exception: traceback.print_exc() - def __download_dataset_threaded(self): - download_thread = threading.Thread(target=self.__download_dataset, daemon=True) + def download_dataset_threaded(self): + download_thread = threading.Thread(target=self.download_dataset, daemon=True) download_thread.start() - def _read_text_file_for_preview(self, file_path: str) -> str: + def _read_text_file_for_preview(self, file_path: str, preview_augmentations: bool) -> str: empty_msg = "[Empty prompt]" try: with open(file_path, "r") as f: - if self.preview_augmentations.get(): + if preview_augmentations: lines = [line.strip() for line in f if line.strip()] return random.choice(lines) if lines else empty_msg content = f.read().strip() @@ -605,7 +85,7 @@ def _read_text_file_for_preview(self, file_path: str) -> str: except UnicodeDecodeError: return "[Invalid file encoding. This should not happen, please report this issue]" - def __get_preview_image(self): + def get_preview_image(self, image_preview_file_index: int, preview_augmentations: bool): preview_image_path = "resources/icons/icon.png" file_index = -1 glob_pattern = "**/*.*" if self.concept.include_subdirectories else "*.*" @@ -620,7 +100,7 @@ def __get_preview_image(self): and not path.name.endswith("-masklabel.png") and not path.name.endswith("-condlabel.png"): preview_image_path = path_util.canonical_join(concept_path, path) file_index += 1 - if file_index == self.image_preview_file_index: + if file_index == image_preview_file_index: break image = load_image(preview_image_path, 'RGB') @@ -647,10 +127,10 @@ def __get_preview_image(self): "concept": pathlib.Path(self.concept.text.prompt_path) if self.concept.text.prompt_path else None, } file_path = file_map.get(source) - prompt_output = self._read_text_file_for_preview(str(file_path)) if file_path else "[Empty prompt]" + prompt_output = self._read_text_file_for_preview(str(file_path), preview_augmentations) if file_path else "[Empty prompt]" modules = [] - if self.preview_augmentations.get(): + if preview_augmentations: input_module = InputPipelineModule({ 'true': True, 'image': image_tensor, @@ -749,133 +229,16 @@ def __get_preview_image(self): return image, filename_output, prompt_output - def __update_concept_stats(self): - #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") - - #directory count - self.dir_count_preview.configure(text=self.concept.concept_stats["directory_count"]) - - #image count - self.image_count_preview.configure(text=self.concept.concept_stats["image_count"]) - self.image_count_mask_preview.configure(text=self.concept.concept_stats["image_with_mask_count"]) - self.image_count_caption_preview.configure(text=self.concept.concept_stats["image_with_caption_count"]) - - #video count - self.video_count_preview.configure(text=self.concept.concept_stats["video_count"]) - #self.video_count_mask_preview.configure(text=self.concept.concept_stats["video_with_mask_count"]) - self.video_count_caption_preview.configure(text=self.concept.concept_stats["video_with_caption_count"]) - - #mask count - self.mask_count_preview.configure(text=self.concept.concept_stats["mask_count"]) - self.mask_count_preview_unpaired.configure(text=self.concept.concept_stats["unpaired_masks"]) - - #caption count - if self.concept.concept_stats["subcaption_count"] > 0: - self.caption_count_preview.configure(text=f'{self.concept.concept_stats["caption_count"]} ({self.concept.concept_stats["subcaption_count"]})') - else: - self.caption_count_preview.configure(text=self.concept.concept_stats["caption_count"]) - self.caption_count_preview_unpaired.configure(text=self.concept.concept_stats["unpaired_captions"]) - - #resolution info - max_pixels = self.concept.concept_stats["max_pixels"] - avg_pixels = self.concept.concept_stats["avg_pixels"] - min_pixels = self.concept.concept_stats["min_pixels"] - - if any(isinstance(x, str) for x in [max_pixels, avg_pixels, min_pixels]) or self.concept.concept_stats["image_count"] == 0: #will be str if adv stats were not taken - self.pixel_max_preview.configure(text="-") - self.pixel_avg_preview.configure(text="-") - self.pixel_min_preview.configure(text="-") - else: - #formatted as (#pixels/1000000) MP, width x height, \n filename - self.pixel_max_preview.configure(text=f'{str(round(max_pixels[0]/1000000, 2))} MP, {max_pixels[2]}\n{max_pixels[1]}') - self.pixel_avg_preview.configure(text=f'{str(round(avg_pixels/1000000, 2))} MP, ~{int(math.sqrt(avg_pixels))}w x {int(math.sqrt(avg_pixels))}h') - self.pixel_min_preview.configure(text=f'{str(round(min_pixels[0]/1000000, 2))} MP, {min_pixels[2]}\n{min_pixels[1]}') - - #video length and fps info - max_length = self.concept.concept_stats["max_length"] - avg_length = self.concept.concept_stats["avg_length"] - min_length = self.concept.concept_stats["min_length"] - max_fps = self.concept.concept_stats["max_fps"] - avg_fps = self.concept.concept_stats["avg_fps"] - min_fps = self.concept.concept_stats["min_fps"] - - if any(isinstance(x, str) for x in [max_length, avg_length, min_length]) or self.concept.concept_stats["video_count"] == 0: #will be str if adv stats were not taken - self.length_max_preview.configure(text="-") - self.length_avg_preview.configure(text="-") - self.length_min_preview.configure(text="-") - self.fps_max_preview.configure(text="-") - self.fps_avg_preview.configure(text="-") - self.fps_min_preview.configure(text="-") - else: - #formatted as (#frames) frames \n filename - self.length_max_preview.configure(text=f'{int(max_length[0])} frames\n{max_length[1]}') - self.length_avg_preview.configure(text=f'{int(avg_length)} frames') - self.length_min_preview.configure(text=f'{int(min_length[0])} frames\n{min_length[1]}') - #formatted as (#fps) fps \n filename - self.fps_max_preview.configure(text=f'{int(max_fps[0])} fps\n{max_fps[1]}') - self.fps_avg_preview.configure(text=f'{int(avg_fps)} fps') - self.fps_min_preview.configure(text=f'{int(min_fps[0])} fps\n{min_fps[1]}') - - #caption info - max_caption_length = self.concept.concept_stats["max_caption_length"] - avg_caption_length = self.concept.concept_stats["avg_caption_length"] - min_caption_length = self.concept.concept_stats["min_caption_length"] - - if any(isinstance(x, str) for x in [max_caption_length, avg_caption_length, min_caption_length]) or self.concept.concept_stats["caption_count"] == 0: #will be str if adv stats were not taken - self.caption_max_preview.configure(text="-") - self.caption_avg_preview.configure(text="-") - self.caption_min_preview.configure(text="-") - else: - #formatted as (#chars) chars, (#words) words, \n filename - self.caption_max_preview.configure(text=f'{max_caption_length[0]} chars, {max_caption_length[2]} words\n{max_caption_length[1]}') - self.caption_avg_preview.configure(text=f'{int(avg_caption_length[0])} chars, {int(avg_caption_length[1])} words') - self.caption_min_preview.configure(text=f'{min_caption_length[0]} chars, {min_caption_length[2]} words\n{min_caption_length[1]}') - - #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) - - self.bucket_ax.cla() - aspects = [str(x) for x in list(aspect_buckets.keys())] - aspect_ratios = [self.decimal_to_aspect_ratio(x) for x in list(aspect_buckets.keys())] - counts = list(aspect_buckets.values()) - b = self.bucket_ax.bar(aspect_ratios, counts) - self.bucket_ax.bar_label(b, color=self.text_color) - sec = self.bucket_ax.secondary_xaxis(location=-0.1) - sec.spines["bottom"].set_linewidth(0) - sec.set_xticks([0, (len(aspects)-1)/2, len(aspects)-1], labels=["Wide", "Square", "Tall"]) - sec.tick_params('x', length=0) - self.canvas.draw() - - def decimal_to_aspect_ratio(self, value : float): - #find closest fraction to decimal aspect value and convert to a:b format - aspect_fraction = fractions.Fraction(value).limit_denominator(16) - aspect_string = f'{aspect_fraction.denominator}:{aspect_fraction.numerator}' - return aspect_string - - def __get_concept_stats(self, advanced_checks: bool, wait_time: float): + def get_concept_stats(self, view, advanced_checks: bool, wait_time: float): start_time = time.perf_counter() last_update = time.perf_counter() self.cancel_scan_flag.clear() - self.concept_stats_tab.after(0, self.__disable_scan_buttons) + view.components.call_after(view.concept_stats_tab, 0, view._disable_scan_buttons) concept_path = self.get_concept_path(self.concept.path) if not concept_path: print(f"Unable to get statistics for concept path: {self.concept.path}") - self.concept_stats_tab.after(0, self.__enable_scan_buttons) + view.components.call_after(view.concept_stats_tab, 0, view._enable_scan_buttons) return subfolders = [concept_path] @@ -890,45 +253,45 @@ def __get_concept_stats(self, advanced_checks: bool, wait_time: float): #update GUI approx every half second if time.perf_counter() > (last_update + 0.5): last_update = time.perf_counter() - self.concept_stats_tab.after(0, self.__update_concept_stats) + view.components.call_after(view.concept_stats_tab, 0, lambda: view._update_concept_stats(self)) 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) + view.components.call_after(view.concept_stats_tab, 0, view._enable_scan_buttons) + view.components.call_after(view.concept_stats_tab, 0, lambda: view._update_concept_stats(self)) - 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) + def get_concept_stats_threaded(self, view, advanced_checks: bool, waittime: float): + self.scan_thread = threading.Thread(target=self.get_concept_stats, args=[view, advanced_checks, waittime], daemon=True) self.scan_thread.start() - def __disable_scan_buttons(self): - self.refresh_basic_stats_button.configure(state="disabled") - self.refresh_advanced_stats_button.configure(state="disabled") - - def __enable_scan_buttons(self): - self.refresh_basic_stats_button.configure(state="normal") - self.refresh_advanced_stats_button.configure(state="normal") - - def __cancel_concept_stats(self): - self.cancel_scan_flag.set() - - def __auto_update_concept_stats(self): + def auto_update_concept_stats(self, view): try: - self.__update_concept_stats() #load stats from config if available, else raises KeyError + view._update_concept_stats(self) #load stats from config if available, else raises KeyError if self.concept.concept_stats["file_size"] == 0: #force rescan if empty raise KeyError except KeyError: concept_path = self.get_concept_path(self.concept.path) if concept_path: - self.__get_concept_stats(False, 2) #force rescan if config is empty, timeout of 2 sec + self.get_concept_stats(view, 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 + self.get_concept_stats(view, True, 2) #do advanced scan automatically if basic took <0.1s + + +class InputPipelineModule( + PipelineModule, + RandomAccessPipelineModule, +): + def __init__(self, data: dict): + super().__init__() + self.data = data - def destroy(self): - if self.bucket_fig is not None: - plt.close(self.bucket_fig) - self.bucket_fig = None + def length(self) -> int: + return 1 - super().destroy() + def get_inputs(self) -> list[str]: + return [] - def __ok(self): - self.destroy() + def get_outputs(self) -> list[str]: + return list(self.data.keys()) + + def get_item(self, variation: int, index: int, requested_name: str = None) -> dict: + return self.data diff --git a/modules/ui/ConvertModelUIController.py b/modules/ui/ConvertModelUIController.py index 6cb1b507a..c3cdebd31 100644 --- a/modules/ui/ConvertModelUIController.py +++ b/modules/ui/ConvertModelUIController.py @@ -4,123 +4,23 @@ from modules.util import create from modules.util.args.ConvertModelArgs import ConvertModelArgs from modules.util.config.TrainConfig import QuantizationConfig -from modules.util.enum.DataType import DataType -from modules.util.enum.ModelFormat import ModelFormat -from modules.util.enum.ModelType import ModelType -from modules.util.enum.PathIOType import PathIOType from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.ModelNames import EmbeddingName, ModelNames from modules.util.torch_util import torch_gc -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState -import customtkinter as ctk - -class ConvertModelUI(ctk.CTkToplevel): - def __init__(self, parent, *args, **kwargs): - super().__init__(parent, *args, **kwargs) - self.parent = parent - - self.parent = parent +class ConvertModelUIController: + def __init__(self): self.convert_model_args = ConvertModelArgs.default_values() - self.ui_state = UIState(self, self.convert_model_args) - self.button = None - - - self.title("Convert models") - self.geometry("550x350") - self.resizable(True, True) - - self.frame = ctk.CTkFrame(self, width=600, height=300) - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) - - self.main_frame(self.frame) - self.frame.pack(fill="both", expand=True) - - self.wait_visibility() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def main_frame(self, master): - # model type - components.label(master, 0, 0, "Model Type", - tooltip="Type of the model") - components.options_kv(master, 0, 1, [ #TODO simplify - ("Stable Diffusion 1.5", ModelType.STABLE_DIFFUSION_15), - ("Stable Diffusion 1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), - ("Stable Diffusion 2.0", ModelType.STABLE_DIFFUSION_20), - ("Stable Diffusion 2.0 Inpainting", ModelType.STABLE_DIFFUSION_20_INPAINTING), - ("Stable Diffusion 2.1", ModelType.STABLE_DIFFUSION_21), - ("Stable Diffusion 3", ModelType.STABLE_DIFFUSION_3), - ("Stable Diffusion 3.5", ModelType.STABLE_DIFFUSION_35), - ("Stable Diffusion XL 1.0 Base", ModelType.STABLE_DIFFUSION_XL_10_BASE), - ("Stable Diffusion XL 1.0 Base Inpainting", ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING), - ("Wuerstchen v2", ModelType.WUERSTCHEN_2), - ("Stable Cascade", ModelType.STABLE_CASCADE_1), - ("PixArt Alpha", ModelType.PIXART_ALPHA), - ("PixArt Sigma", ModelType.PIXART_SIGMA), - ("Flux Dev", ModelType.FLUX_DEV_1), - ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), - ("Flux 2", ModelType.FLUX_2), - ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), - ("Chroma1", ModelType.CHROMA_1), #TODO does this just work? HiDream is not here - ("QwenImage", ModelType.QWEN), #TODO does this just work? HiDream is not here - ("ZImage", ModelType.Z_IMAGE), - ], self.ui_state, "model_type") - - # training method - components.label(master, 1, 0, "Model Type", - tooltip="The type of model to convert") - components.options_kv(master, 1, 1, [ - ("Base Model", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ("Embedding", TrainingMethod.EMBEDDING), - ], self.ui_state, "training_method") - - # input name - components.label(master, 2, 0, "Input name", - tooltip="Filename, directory or hugging face repository of the base model") - components.path_entry( - master, 2, 1, self.ui_state, "input_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # output data type - components.label(master, 3, 0, "Output Data Type", - tooltip="Precision to use when saving the output model") - components.options_kv(master, 3, 1, [ - ("float32", DataType.FLOAT_32), - ("float16", DataType.FLOAT_16), - ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "output_dtype") - - # output format - components.label(master, 4, 0, "Output Format", - tooltip="Format to use when saving the output model") - components.options_kv(master, 4, 1, [ - ("Safetensors", ModelFormat.SAFETENSORS), - ("Diffusers", ModelFormat.DIFFUSERS), - ], self.ui_state, "output_model_format") - - # output model destination - components.label(master, 5, 0, "Model Output Destination", - tooltip="Filename or directory where the output model is saved") - components.path_entry( - master, 5, 1, self.ui_state, "output_model_destination", - mode="file", - io_type=PathIOType.MODEL, - ) + self.view = None - self.button = components.button(master, 6, 1, "Convert", self.convert_model) + def create_window(self, parent, view_cls): + self.view = view_cls(parent, self) + return self.view def convert_model(self): try: - self.button.configure(state="disabled") + self.view.set_converting(True) model_loader = create.create_model_loader( model_type=self.convert_model_args.model_type, training_method=self.convert_model_args.training_method @@ -167,4 +67,4 @@ def convert_model(self): traceback.print_exc() torch_gc() - self.button.configure(state="normal") + self.view.set_converting(False) diff --git a/modules/ui/CtkAdditionalEmbeddingsTabView.py b/modules/ui/CtkAdditionalEmbeddingsTabView.py index 6a5e3fbe7..fc24c61d1 100644 --- a/modules/ui/CtkAdditionalEmbeddingsTabView.py +++ b/modules/ui/CtkAdditionalEmbeddingsTabView.py @@ -1,54 +1,38 @@ -from modules.ui.ConfigList import ConfigList -from modules.util.config.TrainConfig import TrainConfig, TrainEmbeddingConfig -from modules.util.ui import components -from modules.util.ui.UIState import UIState +from modules.ui.AdditionalEmbeddingsTabController import AdditionalEmbeddingsTabController +from modules.ui.BaseAdditionalEmbeddingsTabView import BaseAdditionalEmbeddingsTabView, BaseEmbeddingWidgetView +from modules.ui.CtkConfigListView import CtkConfigListView +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState import customtkinter as ctk -class AdditionalEmbeddingsTab(ConfigList): +class CtkAdditionalEmbeddingsTabView(CtkConfigListView, BaseAdditionalEmbeddingsTabView): - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__( - master, - train_config, - ui_state, + def __init__(self, master, controller: AdditionalEmbeddingsTabController, ui_state): + CtkConfigListView.__init__( + self, master, controller, ui_state, attr_name="additional_embeddings", enable_key="train", from_external_file=False, add_button_text="add embedding", is_full_width=True, - show_toggle_button=True + show_toggle_button=True, ) - def refresh_ui(self): - if self.element_list is not None: - self.element_list.destroy() - self.element_list = None - self.widgets_initialized = False - self._create_element_list() - def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - return EmbeddingWidget(master, element, i, open_command, remove_command, clone_command, save_command) - - def create_new_element(self) -> dict: - return TrainEmbeddingConfig.default_values() + return CtkEmbeddingWidgetView(master, element, i, open_command, remove_command, clone_command, save_command, self.controller) - def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: - pass +class CtkEmbeddingWidgetView(BaseEmbeddingWidgetView, ctk.CTkFrame): -class EmbeddingWidget(ctk.CTkFrame): - def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): - super().__init__( - master=master, corner_radius=10, bg_color="transparent" - ) + def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command, controller): + ctk.CTkFrame.__init__(self, master=master, corner_radius=10, bg_color="transparent") + BaseEmbeddingWidgetView.__init__(self, ctk_components) self.element = element - self.ui_state = UIState(self, element) - self.i = i - self.save_command = save_command + ui_state = CtkUIState(self, element) self.grid_columnconfigure(0, weight=1) @@ -61,76 +45,7 @@ def __init__(self, master, element, i, open_command, remove_command, clone_comma bottom_frame.grid(row=1, column=0, sticky="nsew") bottom_frame.grid_columnconfigure(7, weight=1) - # close button - close_button = ctk.CTkButton( - master=top_frame, - width=20, - height=20, - text="X", - corner_radius=2, - fg_color="#C00000", - command=lambda: remove_command(self.i), - ) - close_button.grid(row=0, column=0) - - # clone button - clone_button = ctk.CTkButton( - master=top_frame, - width=20, - height=20, - text="+", - corner_radius=2, - fg_color="#00C000", - command=lambda: clone_command(self.i, self.__randomize_uuid), - ) - clone_button.grid(row=0, column=1, padx=5) - - # 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.path_entry( - top_frame, 0, 3, self.ui_state, "model_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # placeholder - components.label(top_frame, 0, 4, "placeholder:", - tooltip="The placeholder used when using the embedding in a prompt") - components.entry(top_frame, 0, 5, self.ui_state, "placeholder") - - # token count - components.label(top_frame, 0, 6, "token count:", - tooltip="The token count used when creating a new embedding. Leave empty to auto detect from the initial embedding text.") - token_count_entry = components.entry(top_frame, 0, 7, self.ui_state, "token_count") - token_count_entry.configure(width=40) - - # trainable - components.label(bottom_frame, 0, 0, "train:") - trainable_switch = components.switch(bottom_frame, 0, 1, self.ui_state, "train", command=save_command) - trainable_switch.configure(width=40) - - # output embedding - components.label(bottom_frame, 0, 2, "output embedding:", - tooltip="Output embeddings are calculated at the output of the text encoder, not the input. This can improve results for larger text encoders and lower VRAM usage.") - output_embedding_switch = components.switch(bottom_frame, 0, 3, self.ui_state, "is_output_embedding") - output_embedding_switch.configure(width=40) - - # stop training after - components.label(bottom_frame, 0, 4, "stop training after:", - tooltip="When to stop training the embedding") - components.time_entry(bottom_frame, 0, 5, self.ui_state, "stop_training_after", "stop_training_after_unit") - - # initial embedding text - components.label(bottom_frame, 0, 6, "initial embedding text:", - tooltip="The initial embedding text used when creating a new embedding") - components.entry(bottom_frame, 0, 7, self.ui_state, "initial_embedding_text") - - def __randomize_uuid(self, embedding_config: TrainEmbeddingConfig): - embedding_config.uuid = TrainEmbeddingConfig.default_values().uuid - return embedding_config - - def configure_element(self): - pass + self.build_content(top_frame, bottom_frame, ui_state, i, save_command, remove_command, clone_command, controller) def place_in_list(self): self.grid(row=self.i, column=0, pady=5, padx=5, sticky="new") diff --git a/modules/ui/CtkCaptionUIView.py b/modules/ui/CtkCaptionUIView.py index e6cc0551e..281036912 100644 --- a/modules/ui/CtkCaptionUIView.py +++ b/modules/ui/CtkCaptionUIView.py @@ -1,97 +1,47 @@ -import os -import platform -import subprocess -import traceback from tkinter import filedialog -from modules.module.Blip2Model import Blip2Model -from modules.module.BlipModel import BlipModel -from modules.module.ClipSegModel import ClipSegModel -from modules.module.MaskByColor import MaskByColor -from modules.module.RembgHumanModel import RembgHumanModel -from modules.module.RembgModel import RembgModel -from modules.module.WDModel import WDModel -from modules.ui.GenerateCaptionsWindow import GenerateCaptionsWindow -from modules.ui.GenerateMasksWindow import GenerateMasksWindow -from modules.util import path_util -from modules.util.image_util import load_image -from modules.util.torch_util import default_device, torch_gc -from modules.util.ui import components +from modules.ui.BaseCaptionUIView import BaseCaptionUIView +from modules.ui.CaptionUIController import CaptionUIController +from modules.ui.CtkGenerateCaptionsWindowView import CtkGenerateCaptionsWindowView +from modules.ui.CtkGenerateMasksWindowView import CtkGenerateMasksWindowView +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState from modules.util.ui.ui_utils import bind_mousewheel, set_window_icon -from modules.util.ui.UIState import UIState - -import torch import customtkinter as ctk -import cv2 -import numpy as np from customtkinter import ScalingTracker, ThemeManager -from PIL import Image, ImageDraw - - -class CaptionUI(ctk.CTkToplevel): - def __init__( - self, - parent, - initial_dir: str | None, - initial_include_subdirectories: bool, - *args, - **kwargs, - ) -> None: - super().__init__(parent, *args, **kwargs) - self.protocol("WM_DELETE_WINDOW", self._on_close) - - self.dir = initial_dir - self.config_ui_data = {"include_subdirectories": initial_include_subdirectories} - self.config_ui_state = UIState(self, self.config_ui_data) - self.image_size = 850 - self.help_text = """ - Keyboard shortcuts when focusing on the prompt input field: - Up arrow: previous image - Down arrow: next image - Return: save - Ctrl+M: only show the mask - Ctrl+D: draw mask editing mode - Ctrl+F: fill mask editing mode - - When editing masks: - Left click: add mask - Right click: remove mask - Mouse wheel: increase or decrease brush size""" - self.masking_model = None - self.captioning_model = None - self.image_rel_paths = [] - self.current_image_index = -1 - self.file_list = None - self.image_labels = [] - self.pil_image = None - self.image_width = 0 - self.image_height = 0 - self.pil_mask = None - self.mask_draw_x = 0 - self.mask_draw_y = 0 - self.mask_draw_radius = 0.01 - self.display_only_mask = False - self.image = None - self.image_label = None - self.mask_editing_mode = 'draw' +from PIL import Image + + +class CtkCaptionUIView(BaseCaptionUIView, ctk.CTkToplevel): + def __init__(self, parent, controller: CaptionUIController, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseCaptionUIView.__init__(self, ctk_components) + self.protocol("WM_DELETE_WINDOW", controller.on_close) + + self.controller = controller + controller.view = self + self.config_ui_state = CtkUIState(self, controller.config_ui_data) self.enable_mask_editing_var = ctk.BooleanVar() self.mask_editing_alpha = None self.prompt_var = None self.prompt_component = None - + self.image = None + self.image_label = None + self.file_list = None + self.image_labels = [] self.title("OneTrainer") self.geometry("1280x980") self.resizable(False, False) - self.grid_rowconfigure(0, weight=0) self.grid_rowconfigure(1, weight=1) self.grid_columnconfigure(0, weight=1) - - self.top_bar(self) + top_frame = ctk.CTkFrame(self) + top_frame.grid(row=0, column=0, sticky="nsew") + self.build_top_bar(top_frame, controller, self.config_ui_state) self.bottom_frame = ctk.CTkFrame(self) self.bottom_frame.grid(row=1, column=0, sticky="nsew") @@ -101,35 +51,12 @@ def __init__( self.file_list_column(self.bottom_frame) self.content_column(self.bottom_frame) - self.load_directory() + self.controller.load_directory() self.wait_visibility() self.focus_set() self.after(200, lambda: set_window_icon(self)) - def top_bar(self, master): - top_frame = ctk.CTkFrame(master) - top_frame.grid(row=0, column=0, sticky="nsew") - - components.button(top_frame, 0, 0, "Open", self.open_directory, - tooltip="open a new directory") - components.button(top_frame, 0, 1, "Generate Masks", self.open_mask_window, - tooltip="open a dialog to automatically generate masks") - components.button(top_frame, 0, 2, "Generate Captions", self.open_caption_window, - tooltip="open a dialog to automatically generate captions") - - if platform.system() == "Windows": - components.button(top_frame, 0, 3, "Open in Explorer", self.open_in_explorer, - tooltip="open the current image in Explorer") - - components.switch(top_frame, 0, 4, self.config_ui_state, "include_subdirectories", - text="include subdirectories") - - top_frame.grid_columnconfigure(5, weight=1) - - components.button(top_frame, 0, 6, "Help", self.print_help, - tooltip=self.help_text) - def file_list_column(self, master): if self.file_list is not None: self.image_labels = [] @@ -138,10 +65,10 @@ def file_list_column(self, master): self.file_list = ctk.CTkScrollableFrame(master, width=300) self.file_list.grid(row=0, column=0, sticky="nsew") - for i, filename in enumerate(self.image_rel_paths): + for i, filename in enumerate(self.controller.image_rel_paths): def __create_switch_image(index): def __switch_image(event): - self.switch_image(index) + self.controller.switch_image(index) return __switch_image @@ -160,10 +87,7 @@ def content_column(self, master): right_frame.grid_columnconfigure(4, weight=1) right_frame.grid_rowconfigure(1, weight=1) - components.button(right_frame, 0, 0, "Draw", self.draw_mask_editing_mode, - tooltip="draw a mask using a brush") - components.button(right_frame, 0, 1, "Fill", self.fill_mask_editing_mode, - tooltip="draw a mask using a fill tool") + self.build_mask_buttons(right_frame) # checkbox to enable mask editing self.enable_mask_editing_var = ctk.BooleanVar() @@ -184,10 +108,11 @@ def content_column(self, master): # image self.image = ctk.CTkImage( light_image=image, - size=(self.image_size, self.image_size) + size=(self.controller.image_size, self.controller.image_size) ) self.image_label = ctk.CTkLabel( - master=right_frame, text="", image=self.image, height=self.image_size, width=self.image_size + master=right_frame, text="", image=self.image, + height=self.controller.image_size, width=self.controller.image_size ) self.image_label.grid(row=1, column=0, columnspan=5, sticky="nsew") @@ -204,156 +129,37 @@ def content_column(self, master): self.prompt_component.focus_set() def bind_key_events(self, component): - component.bind("", self.next_image) - component.bind("", self.previous_image) + component.bind("", lambda e: self.controller.next_image()) + component.bind("", lambda e: self.controller.previous_image()) component.bind("", self.save) component.bind("", self.toggle_mask) component.bind("", self.draw_mask_editing_mode) component.bind("", self.fill_mask_editing_mode) - def load_directory(self, include_subdirectories: bool = False): - self.scan_directory(include_subdirectories) + def refresh_file_list(self): self.file_list_column(self.bottom_frame) - if len(self.image_rel_paths) > 0: - self.switch_image(0) - else: - self.switch_image(-1) - + def focus_prompt(self): self.prompt_component.focus_set() - def scan_directory(self, include_subdirectories: bool = False): - def __is_supported_image_extension(filename): - name, ext = os.path.splitext(filename) - return path_util.is_supported_image_extension(ext) and not name.endswith("-masklabel") and not name.endswith("-condlabel") - - self.image_rel_paths = [] - - if not self.dir or not os.path.isdir(self.dir): - return - - if include_subdirectories: - for root, _, files in os.walk(self.dir): - for filename in files: - if __is_supported_image_extension(filename): - self.image_rel_paths.append( - os.path.relpath(os.path.join(root, filename), self.dir) - ) - else: - for _, filename in enumerate(os.listdir(self.dir)): - if __is_supported_image_extension(filename): - self.image_rel_paths.append( - os.path.relpath(os.path.join(self.dir, filename), self.dir) - ) - - def load_image(self): - image_name = "resources/icons/icon.png" - - if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): - image_name = self.image_rel_paths[self.current_image_index] - image_name = os.path.join(self.dir, image_name) - - try: - return load_image(image_name, convert_mode="RGB") - except Exception: - print(f'Could not open image {image_name}') - - def load_mask(self): - if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): - image_name = self.image_rel_paths[self.current_image_index] - mask_name = os.path.splitext(image_name)[0] + "-masklabel.png" - mask_name = os.path.join(self.dir, mask_name) - - try: - return load_image(mask_name, convert_mode='RGB') - except Exception: - return None - else: - return None - - def load_prompt(self): - if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): - image_name = self.image_rel_paths[self.current_image_index] - prompt_name = os.path.splitext(image_name)[0] + ".txt" - prompt_name = os.path.join(self.dir, prompt_name) - - try: - with open(prompt_name, "r", encoding='utf-8') as f: - return f.readlines()[0].strip() - except Exception: - return "" - else: - return "" - - def previous_image(self, event): - if len(self.image_rel_paths) > 0 and (self.current_image_index - 1) >= 0: - self.switch_image(self.current_image_index - 1) - - def next_image(self, event): - if len(self.image_rel_paths) > 0 and (self.current_image_index + 1) < len(self.image_rel_paths): - self.switch_image(self.current_image_index + 1) - - def switch_image(self, index): - if len(self.image_labels) > 0 and self.current_image_index < len(self.image_labels): - self.image_labels[self.current_image_index].configure( + def on_image_switched(self, old_index, new_index, prompt): + if len(self.image_labels) > 0 and old_index < len(self.image_labels): + self.image_labels[old_index].configure( text_color=ThemeManager.theme["CTkLabel"]["text_color"]) + self.image_labels[new_index].configure(text_color="#FF0000") + self.refresh_image() + self.prompt_var.set(prompt) - self.current_image_index = index - if index >= 0: - self.image_labels[index].configure(text_color="#FF0000") - - self.pil_image = self.load_image() - self.pil_mask = self.load_mask() - prompt = self.load_prompt() - - self.image_width = self.pil_image.width - self.image_height = self.pil_image.height - scale = self.image_size / max(self.pil_image.height, self.pil_image.width) - height = int(self.pil_image.height * scale) - width = int(self.pil_image.width * scale) - - self.pil_image = self.pil_image.resize((width, height), Image.Resampling.LANCZOS) - - self.refresh_image() - self.prompt_var.set(prompt) - else: - image = Image.new("RGB", (512, 512), (0, 0, 0)) - self.image.configure(light_image=image) + def on_image_cleared(self): + image = Image.new("RGB", (512, 512), (0, 0, 0)) + self.image.configure(light_image=image) def refresh_image(self): - if self.pil_mask: - resized_pil_mask = self.pil_mask.resize( - (self.pil_image.width, self.pil_image.height), - Image.Resampling.NEAREST - ) - - if self.display_only_mask: - self.image.configure(light_image=resized_pil_mask, size=resized_pil_mask.size) - else: - np_image = np.array(self.pil_image).astype(np.float32) / 255.0 - np_mask = np.array(resized_pil_mask).astype(np.float32) / 255.0 - - # normalize mask between 0.3 - 1.0 so we can see image underneath and gauge strength of the alpha - norm_min = 0.3 - np_mask_min = np_mask.min() - if np_mask_min == 0: - # optimize for common case - np_mask = np_mask * (1.0 - norm_min) + norm_min - elif np_mask_min < 1: - # note: min of 1 means we get divide by 0 - np_mask = (np_mask - np_mask_min) / (1.0 - np_mask_min) * (1.0 - norm_min) + norm_min - - np_masked_image = (np_image * np_mask * 255.0).astype(np.uint8) - masked_image = Image.fromarray(np_masked_image, mode='RGB') - - self.image.configure(light_image=masked_image, size=masked_image.size) - else: - self.image.configure(light_image=self.pil_image, size=self.pil_image.size) + pil_image, size = self.controller.get_display_image() + self.image.configure(light_image=pil_image, size=size) def draw_mask_radius(self, delta, raw_event): - # Wheel up = Increase radius. Wheel down = Decrease radius. - multiplier = 1.0 + (delta * 0.05) - self.mask_draw_radius = max(0.0025, self.mask_draw_radius * multiplier) + self.controller.update_mask_draw_radius(delta) def edit_mask(self, event): if not self.enable_mask_editing_var.get(): @@ -362,22 +168,11 @@ def edit_mask(self, event): if event.widget != self.image_label.children["!label"]: return - if len(self.image_rel_paths) == 0 or self.current_image_index >= len(self.image_rel_paths): - return - display_scaling = ScalingTracker.get_window_scaling(self) event_x = event.x / display_scaling event_y = event.y / display_scaling - start_x = int(event_x / self.pil_image.width * self.image_width) - start_y = int(event_y / self.pil_image.height * self.image_height) - end_x = int(self.mask_draw_x / self.pil_image.width * self.image_width) - end_y = int(self.mask_draw_y / self.pil_image.height * self.image_height) - - self.mask_draw_x = event_x - self.mask_draw_y = event_y - is_right = False is_left = False if event.state & 0x0100 or event.num == 1: # left mouse button @@ -385,96 +180,18 @@ def edit_mask(self, event): elif event.state & 0x0400 or event.num == 3: # right mouse button is_right = True - if self.mask_editing_mode == 'draw': - self.draw_mask(start_x, start_y, end_x, end_y, is_left, is_right) - if self.mask_editing_mode == 'fill': - self.fill_mask(start_x, start_y, end_x, end_y, is_left, is_right) - - def draw_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): - color = None - - adding_to_mask = True - if is_left: - try: - alpha = float(self.mask_editing_alpha.get()) - except Exception: - alpha = 1.0 - rgb_value = int(max(0.0, min(alpha, 1.0)) * 255) # max/min stuff to clamp to 0 - 255 range - color = (rgb_value, rgb_value, rgb_value) - - elif is_right: - color = (0, 0, 0) - adding_to_mask = False - - if color is not None: - if self.pil_mask is None: - if adding_to_mask: - self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(0, 0, 0)) - else: - self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(255, 255, 255)) - - radius = int(self.mask_draw_radius * max(self.pil_mask.width, self.pil_mask.height)) - - draw = ImageDraw.Draw(self.pil_mask) - draw.line((start_x, start_y, end_x, end_y), fill=color, - width=radius + radius + 1) - draw.ellipse((start_x - radius, start_y - radius, - start_x + radius, start_y + radius), fill=color, outline=None) - draw.ellipse((end_x - radius, end_y - radius, end_x + radius, - end_y + radius), fill=color, outline=None) - - self.refresh_image() - - def fill_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): - color = None - - adding_to_mask = True - if is_left: - try: - alpha = float(self.mask_editing_alpha.get()) - except Exception: - alpha = 1.0 - rgb_value = int(max(0.0, min(alpha, 1.0)) * 255) # max/min stuff to clamp to 0 - 255 range - color = (rgb_value, rgb_value, rgb_value) - - elif is_right: - color = (0, 0, 0) - adding_to_mask = False - - if color is not None: - if self.pil_mask is None: - if adding_to_mask: - self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(0, 0, 0)) - else: - self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(255, 255, 255)) - - np_mask = np.array(self.pil_mask).astype(np.uint8) - cv2.floodFill(np_mask, None, (start_x, start_y), color) - self.pil_mask = Image.fromarray(np_mask, 'RGB') - - self.refresh_image() - - def save(self, event): - if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): - image_name = self.image_rel_paths[self.current_image_index] - - prompt_name = os.path.splitext(image_name)[0] + ".txt" - prompt_name = os.path.join(self.dir, prompt_name) - - mask_name = os.path.splitext(image_name)[0] + "-masklabel.png" - mask_name = os.path.join(self.dir, mask_name) + try: + alpha = float(self.mask_editing_alpha.get()) + except Exception: + alpha = 1.0 - try: - with open(prompt_name, "w", encoding='utf-8') as f: - f.write(self.prompt_var.get()) - except Exception: - return + self.controller.handle_edit_mask(event_x, event_y, is_left, is_right, alpha) - if self.pil_mask: - self.pil_mask.save(mask_name) + def save(self, event): + self.controller.save(self.prompt_var.get()) def draw_mask_editing_mode(self, *args): - self.mask_editing_mode = 'draw' + self.controller.set_mask_editing_mode('draw') if args: # disable default event @@ -482,91 +199,30 @@ def draw_mask_editing_mode(self, *args): return None def fill_mask_editing_mode(self, *args): - self.mask_editing_mode = 'fill' + self.controller.set_mask_editing_mode('fill') def toggle_mask(self, *args): - self.display_only_mask = not self.display_only_mask + self.controller.toggle_mask() self.refresh_image() def open_directory(self): new_dir = filedialog.askdirectory() if new_dir: - self.dir = new_dir - self.load_directory(include_subdirectories=self.config_ui_data["include_subdirectories"]) + self.controller.dir = new_dir + self.controller.load_directory(include_subdirectories=self.controller.config_ui_data["include_subdirectories"]) def open_mask_window(self): - dialog = GenerateMasksWindow(self, self.dir, self.config_ui_data["include_subdirectories"]) - self.wait_window(dialog) - self.switch_image(self.current_image_index) + self.wait_window(self.controller.open_mask_window(self, CtkGenerateMasksWindowView)) + self.controller.switch_image(self.controller.current_image_index) def open_caption_window(self): - dialog = GenerateCaptionsWindow(self, self.dir, self.config_ui_data["include_subdirectories"]) - self.wait_window(dialog) - self.switch_image(self.current_image_index) + self.wait_window(self.controller.open_caption_window(self, CtkGenerateCaptionsWindowView)) + self.controller.switch_image(self.controller.current_image_index) def open_in_explorer(self): - try: - image_name = self.image_rel_paths[self.current_image_index] - image_name = os.path.realpath(os.path.join(self.dir, image_name)) - subprocess.Popen(f"explorer /select,{image_name}") - except Exception: - traceback.print_exc() - - def load_masking_model(self, model): - model_type = type(self.masking_model).__name__ if self.masking_model else None - - if model == "ClipSeg" and model_type != "ClipSegModel": - self._release_models() - print("loading ClipSeg model, this may take a while") - self.masking_model = ClipSegModel(default_device, torch.float32) - elif model == "Rembg" and model_type != "RembgModel": - self._release_models() - print("loading Rembg model, this may take a while") - self.masking_model = RembgModel(default_device, torch.float32) - elif model == "Rembg-Human" and model_type != "RembgHumanModel": - self._release_models() - print("loading Rembg-Human model, this may take a while") - self.masking_model = RembgHumanModel(default_device, torch.float32) - elif model == "Hex Color" and model_type != "MaskByColor": - self._release_models() - self.masking_model = MaskByColor(default_device, torch.float32) - - def load_captioning_model(self, model): - model_type = type(self.captioning_model).__name__ if self.captioning_model else None - - if model == "Blip" and model_type != "BlipModel": - self._release_models() - print("loading Blip model, this may take a while") - self.captioning_model = BlipModel(default_device, torch.float16) - elif model == "Blip2" and model_type != "Blip2Model": - self._release_models() - print("loading Blip2 model, this may take a while") - self.captioning_model = Blip2Model(default_device, torch.float16) - elif model == "WD14 VIT v2" and model_type != "WDModel": - self._release_models() - print("loading WD14_VIT_v2 model, this may take a while") - self.captioning_model = WDModel(default_device, torch.float16) - - def print_help(self): - print(self.help_text) - - def _release_models(self): - """Release all models from VRAM""" - freed = False - if self.captioning_model is not None: - self.captioning_model = None - freed = True - if self.masking_model is not None: - self.masking_model = None - freed = True - if freed: - torch_gc() - - def _on_close(self): - self._release_models() - self.destroy() + self.controller.open_in_explorer() def destroy(self): - self._release_models() + self.controller._release_models() super().destroy() diff --git a/modules/ui/CtkCloudTabView.py b/modules/ui/CtkCloudTabView.py index 99057e428..0a5249069 100644 --- a/modules/ui/CtkCloudTabView.py +++ b/modules/ui/CtkCloudTabView.py @@ -1,26 +1,18 @@ -import webbrowser -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.CloudAction import CloudAction -from modules.util.enum.CloudFileSync import CloudFileSync -from modules.util.enum.CloudType import CloudType -from modules.util.ui import components -from modules.util.ui.UIState import UIState +from modules.ui.BaseCloudTabView import BaseCloudTabView +from modules.ui.CloudTabController import CloudTabController +from modules.util.ui import ctk_components import customtkinter as ctk -class CloudTab: - - def __init__(self, master, train_config: TrainConfig, ui_state: UIState, parent): - super().__init__() - +class CtkCloudTabView(BaseCloudTabView): + def __init__(self, master, controller: CloudTabController, ui_state): + BaseCloudTabView.__init__(self, ctk_components) self.master = master - self.train_config = train_config + self.controller = controller self.ui_state = ui_state - self.parent = parent - self.reattach = False self.frame = ctk.CTkScrollableFrame(master, fg_color="transparent") self.frame.grid_columnconfigure(0, weight=0) @@ -30,192 +22,24 @@ def __init__(self, master, train_config: TrainConfig, ui_state: UIState, parent) self.frame.grid_columnconfigure(4, weight=0) self.frame.grid_columnconfigure(5, weight=1) - components.label(self.frame, 0, 0, "Enabled", - tooltip="Enable cloud training") - components.switch(self.frame, 0, 1, self.ui_state, "cloud.enabled") - - components.label(self.frame, 1, 0, "Type", - tooltip="Choose LINUX to connect to a linux machine via SSH. Choose RUNPOD for additional functionality such as automatically creating and deleting pods.") - components.options_kv(self.frame, 1, 1, [ - ("RUNPOD", CloudType.RUNPOD), - ("LINUX", CloudType.LINUX), - ], self.ui_state, "cloud.type") - - components.label(self.frame, 2, 0, "File sync method", - tooltip="Choose NATIVE_SCP to use scp.exe to transfer files. FABRIC_SFTP uses the Paramiko/Fabric SFTP implementation for file transfers instead.") - components.options_kv(self.frame, 2, 1, [ - ("NATIVE_SCP", CloudFileSync.NATIVE_SCP), - ("FABRIC_SFTP", CloudFileSync.FABRIC_SFTP), - ], self.ui_state, "cloud.file_sync") - - components.label(self.frame, 3, 0, "API key", - tooltip="Cloud service API key for RUNPOD. Leave empty for LINUX. This value is stored separately, not saved to your configuration file. ") - components.entry(self.frame, 3, 1, self.ui_state, "secrets.cloud.api_key") - - components.label(self.frame, 4, 0, "Hostname", - tooltip="SSH server hostname or IP. Leave empty if you have a Cloud ID or want to automatically create a new cloud.") - components.entry(self.frame, 4, 1, self.ui_state, "secrets.cloud.host") - - components.label(self.frame, 5, 0, "Port", - tooltip="SSH server port. Leave empty if you have a Cloud ID or want to automatically create a new cloud.") - components.entry(self.frame, 5, 1, self.ui_state, "secrets.cloud.port") - - components.label(self.frame, 6, 0, "User", - tooltip='SSH username. Use "root" for RUNPOD. Your SSH client must be set up to connect to the cloud using a public key, without a password. For RUNPOD, create an ed25519 key locally, and copy the contents of the public keyfile to your "SSH Public Keys" on the RunPod website.') - components.entry(self.frame, 6, 1, self.ui_state, "secrets.cloud.user") - - 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.path_entry(self.frame, 7, 1, self.ui_state, "secrets.cloud.key_file", mode="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.") - components.entry(self.frame, 8, 1, self.ui_state, "secrets.cloud.password") - - components.label(self.frame, 9, 0, "Cloud id", - tooltip="RUNPOD Cloud ID. The cloud service must have a public IP and SSH service. Leave empty if you want to automatically create a new RUNPOD cloud, or if you're connecting to another cloud provider via SSH Hostname and Port.") - components.entry(self.frame, 9, 1, self.ui_state, "secrets.cloud.id") - - components.label(self.frame, 10, 0, "Tensorboard TCP tunnel", - tooltip="Instead of starting tensorboard locally, make a TCP tunnel to a tensorboard on the cloud") - components.switch(self.frame, 10, 1, self.ui_state, "cloud.tensorboard_tunnel") + self.build_content(self.frame, controller, ui_state) + self.frame.pack(fill="both", expand=1) - components.label(self.frame, 1, 2, "Remote Directory", - tooltip="The directory on the cloud where files will be uploaded and downloaded.") - components.entry(self.frame, 1, 3, self.ui_state, "cloud.remote_dir") - components.label(self.frame, 2, 2, "OneTrainer Directory", - tooltip="The directory for OneTrainer on the cloud.") - components.entry(self.frame, 2, 3, self.ui_state, "cloud.onetrainer_dir") - components.label(self.frame, 3, 2, "Huggingface cache Directory", - tooltip="Huggingface models are downloaded to this remote directory.") - components.entry(self.frame, 3, 3, self.ui_state, "cloud.huggingface_cache_dir") - components.label(self.frame, 4, 2, "Install OneTrainer", - tooltip="Automatically install OneTrainer from GitHub if the directory doesn't already exist.") - components.switch(self.frame, 4, 3, self.ui_state, "cloud.install_onetrainer") - components.label(self.frame, 5, 2, "Install command", - tooltip="The command for installing OneTrainer. Leave the default, unless you want to use a development branch of OneTrainer.") - components.entry(self.frame, 5, 3, self.ui_state, "cloud.install_cmd") - components.label(self.frame, 6, 2, "Update OneTrainer", - tooltip="Update OneTrainer if it already exists on the cloud.") - components.switch(self.frame, 6, 3, self.ui_state, "cloud.update_onetrainer") + def _on_set_gpu_types(self): + self.gpu_types_menu.configure(values=self.controller.get_gpu_types()) - components.label(self.frame, 8, 2, "Detach remote trainer", - tooltip="Allows the trainer to keep running even if your connection to the cloud is lost.") - components.switch(self.frame, 8, 3, self.ui_state, "cloud.detach_trainer") - components.label(self.frame, 9, 2, "Reattach id", - tooltip="An id identifying the remotely running trainer. In case you have lost connection or closed OneTrainer, it will try to reattach to this id instead of starting a new remote trainer.") - reattach_frame = ctk.CTkFrame(self.frame, fg_color="transparent") + def _make_reattach_frame(self, frame): + reattach_frame = ctk.CTkFrame(frame, fg_color="transparent") reattach_frame.grid(row=9, column=3, padx=0, pady=0, sticky="new") reattach_frame.grid_columnconfigure(0, weight=1) reattach_frame.grid_columnconfigure(1, weight=1) - components.entry(reattach_frame, 0, 0, self.ui_state, "cloud.run_id", width=60) - components.button(reattach_frame, 0, 1, "Reattach now", self.__reattach) - - components.label(self.frame, 11, 2, "Download samples", - tooltip="Download samples from the remote workspace directory to your local machine.") - components.switch(self.frame, 11, 3, self.ui_state, "cloud.download_samples") - components.label(self.frame, 12, 2, "Download output model", - tooltip="Download the final model after training. You can disable this if you plan to use an automatically saved checkpoint instead.") - components.switch(self.frame, 12, 3, self.ui_state, "cloud.download_output_model") - components.label(self.frame, 13, 2, "Download saved checkpoints", - tooltip="Download the automatically saved training checkpoints from the remote workspace directory to your local machine.") - components.switch(self.frame, 13, 3, self.ui_state, "cloud.download_saves") - components.label(self.frame, 14, 2, "Download backups", - tooltip="Download backups from the remote workspace directory to your local machine. It's usually not necessary to download them, because as long as the backups are still available on the cloud, the training can be restarted using one of the cloud's backups.") - components.switch(self.frame, 14, 3, self.ui_state, "cloud.download_backups") - components.label(self.frame, 15, 2, "Download tensorboard logs", - tooltip="Download TensorBoard event logs from the remote workspace directory to your local machine. They can then be viewed locally in TensorBoard. It is recommended to disable \"Sample to TensorBoard\" to reduce the event log size.") - components.switch(self.frame, 15, 3, self.ui_state, "cloud.download_tensorboard") - components.label(self.frame, 16, 2, "Delete remote workspace", - tooltip="Delete the workspace directory on the cloud after training has finished successfully and data has been downloaded.") - components.switch(self.frame, 16, 3, self.ui_state, "cloud.delete_workspace") + return reattach_frame - components.label(self.frame, 1, 4, "Create cloud via API", - tooltip="Automatically creates a new cloud instance if both Host:Port and Cloud ID are empty. Currently supported for RUNPOD.") - create_frame = ctk.CTkFrame(self.frame, fg_color="transparent") + def _make_create_frame(self, frame): + create_frame = ctk.CTkFrame(frame, fg_color="transparent") create_frame.grid(row=1, column=5, padx=0, pady=0, sticky="new") create_frame.grid_columnconfigure(0, weight=0) create_frame.grid_columnconfigure(1, weight=1) - components.switch(create_frame, 0, 0, self.ui_state, "cloud.create") - components.button(create_frame, 0, 1, "Create cloud via website", self.__create_cloud) - - components.label(self.frame, 2, 4, "Cloud name", - tooltip="The name of the new cloud instance.") - components.entry(self.frame, 2, 5, self.ui_state, "cloud.name") - components.label(self.frame, 3, 4, "Type", - tooltip="Select the RunPod cloud type. See RunPod's website for details.") - components.options_kv(self.frame, 3, 5, [ - ("", ""), - ("Community", "COMMUNITY"), - ("Secure", "SECURE"), - ], self.ui_state, "cloud.sub_type") - - - components.label(self.frame, 4, 4, "GPU", - tooltip="Select the GPU type. Enter an API key before pressing the button.") - - _,gpu_components=components.options_adv(self.frame, 4, 5, [("")], self.ui_state, "cloud.gpu_type",adv_command=self.__set_gpu_types) - self.gpu_types_menu=gpu_components['component'] - - components.label(self.frame, 5, 4, "Volume size", - tooltip="Set the storage volume size in GB. This volume persists only until the cloud is deleted - not a RunPod network volume") - components.entry(self.frame, 5, 5, self.ui_state, "cloud.volume_size") - - components.label(self.frame, 6, 4, "Min download", - tooltip="Set the minimum download speed of the cloud in Mbps.") - components.entry(self.frame, 6, 5, self.ui_state, "cloud.min_download") - - components.label(self.frame, 8, 4, "Action on finish", - tooltip="What to do when training finishes and the data has been fully downloaded: Stop or delete the cloud, or do nothing.") - components.options_kv(self.frame, 8, 5, [ - ("None", CloudAction.NONE), - ("Stop", CloudAction.STOP), - ("Delete", CloudAction.DELETE), - ], self.ui_state, "cloud.on_finish") - - components.label(self.frame, 9, 4, "Action on error", - tooltip="What to do if training stops due to an error: Stop or delete the cloud, or do nothing. Data may be lost.") - components.options_kv(self.frame, 9, 5, [ - ("None", CloudAction.NONE), - ("Stop", CloudAction.STOP), - ("Delete", CloudAction.DELETE), - ], self.ui_state, "cloud.on_error") - - components.label(self.frame, 10, 4, "Action on detached finish", - tooltip="What to do when training finishes, but the client has been detached and cannot download data. Data may be lost.") - components.options_kv(self.frame, 10, 5, [ - ("None", CloudAction.NONE), - ("Stop", CloudAction.STOP), - ("Delete", CloudAction.DELETE), - ], self.ui_state, "cloud.on_detached_finish") - - components.label(self.frame, 11, 4, "Action on detached error", - tooltip="What to if training stops due to an error, but the client has been detached and cannot download data. Data may be lost.") - components.options_kv(self.frame, 11, 5, [ - ("None", CloudAction.NONE), - ("Stop", CloudAction.STOP), - ("Delete", CloudAction.DELETE), - ], self.ui_state, "cloud.on_detached_error") - - self.frame.pack(fill="both", expand=1) - - def __set_gpu_types(self): - self.gpu_types_menu.configure(values=[]) - if self.train_config.cloud.type == CloudType.RUNPOD: - import runpod - runpod.api_key=self.train_config.secrets.cloud.api_key - gpus=runpod.get_gpus() - self.gpu_types_menu.configure(values=[gpu['id'] for gpu in gpus]) - - def __reattach(self): - self.reattach=True - try: - self.parent.start_training() - finally: - self.reattach=False - - def __create_cloud(self): - if self.train_config.cloud.type == CloudType.RUNPOD: - webbrowser.open("https://www.runpod.io/console/deploy?template=1a33vbssq9&type=gpu", new=0, autoraise=False) + return create_frame diff --git a/modules/ui/CtkConceptTabView.py b/modules/ui/CtkConceptTabView.py index 0b6505694..5b3e86ac9 100644 --- a/modules/ui/CtkConceptTabView.py +++ b/modules/ui/CtkConceptTabView.py @@ -1,33 +1,27 @@ -import os -import pathlib from tkinter import BooleanVar, StringVar -from modules.ui.ConceptWindow import ConceptWindow -from modules.ui.ConfigList import ConfigList -from modules.util import path_util -from modules.util.config.ConceptConfig import ConceptConfig -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.ConceptType import ConceptType -from modules.util.image_util import load_image -from modules.util.ui import components -from modules.util.ui.UIState import UIState -from modules.util.ui.validation import DebounceTimer +from modules.ui.BaseConceptTabView import BaseConceptTabView, BaseConceptWidgetView +from modules.ui.ConceptTabController import ConceptTabController +from modules.ui.CtkConceptWindowView import CtkConceptWindowView +from modules.ui.CtkConfigListView import CtkConfigListView +from modules.util.ui import ctk_components +from modules.util.ui.ctk_validation import DebounceTimer +from modules.util.ui.CtkUIState import CtkUIState import customtkinter as ctk -from PIL import Image -class ConceptTab(ConfigList): +class CtkConceptTabView(CtkConfigListView, BaseConceptTabView): - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): + def __init__(self, master, controller: ConceptTabController, ui_state): + # Pre-initialize before CtkConfigListView.__init__ because _reset_filters is + # called during build() via options_kv's immediate update_var() call. self.search_var = StringVar() self.filter_var = StringVar(value="ALL") self.show_disabled_var = BooleanVar(value=True) - super().__init__( - master, - train_config, - ui_state, + CtkConfigListView.__init__( + self, master, controller, ui_state, from_external_file=True, attr_name="concept_file_name", config_dir="training_concepts", @@ -35,22 +29,18 @@ 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, ) self._toolbar = None self._toolbar_is_wrapped = False self._add_search_bar() - # wrap toolbar if too narrow self.top_frame.bind('', lambda e: self._maybe_reposition_toolbar(e.width)) - def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - return ConceptWidget(master, element, i, open_command, remove_command, clone_command, save_command) - - def create_new_element(self) -> dict: - return ConceptConfig.default_values() - def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: - return ConceptWindow(self.master, self.train_config, self.current_config[i], ui_state[0], ui_state[1], ui_state[2]) + return self.controller.open_element_window(self.master, self.current_config[i], ui_state[0], ui_state[1], ui_state[2], CtkConceptWindowView) + + def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): + return CtkConceptWidgetView(master, element, i, open_command, remove_command, clone_command, save_command, self.controller) def _add_search_bar(self): toolbar = ctk.CTkFrame(self.top_frame, fg_color="transparent") @@ -58,8 +48,7 @@ def _add_search_bar(self): toolbar.grid_columnconfigure(2, weight=1) self._toolbar = toolbar - # Search - ctk.CTkLabel(toolbar, text="Search:").grid(row=0, column=0, padx=(0,5)) + ctk.CTkLabel(toolbar, text="Search:").grid(row=0, column=0, padx=(0, 5)) self.search_var = StringVar() self.search_entry = ctk.CTkEntry(toolbar, textvariable=self.search_var, placeholder_text="Filter...", width=200) @@ -67,77 +56,22 @@ def _add_search_bar(self): self._search_debouncer = DebounceTimer(self.search_entry, 300, lambda: self._update_filters()) self.search_var.trace_add("write", lambda *_: self._search_debouncer.call()) - # Spacer ctk.CTkLabel(toolbar, text="").grid(row=0, column=2, padx=5) - # Type filter - ctk.CTkLabel(toolbar, text="Type:").grid(row=0, column=3, padx=(0,5)) + ctk.CTkLabel(toolbar, text="Type:").grid(row=0, column=3, padx=(0, 5)) self.filter_var = StringVar(value="ALL") - ctk.CTkOptionMenu(toolbar, values=["ALL", "STANDARD", "VALIDATION", "PRIOR_PREDICTION"], + ctk.CTkOptionMenu(toolbar, values=self._FILTER_TYPES, variable=self.filter_var, command=lambda x: self._update_filters(), width=150).grid(row=0, column=4) - # Show disabled checkbox self.show_disabled_var = BooleanVar(value=True) self.show_disabled_checkbox = ctk.CTkCheckBox(toolbar, text="Show Disabled", variable=self.show_disabled_var, command=self._update_filters, width=100) - self.show_disabled_checkbox.grid(row=0, column=5, padx=(10,0)) + self.show_disabled_checkbox.grid(row=0, column=5, padx=(10, 0)) self._refresh_show_disabled_text() - # Clear button ctk.CTkButton(toolbar, text="Clear", width=50, - command=self._reset_filters).grid(row=0, column=6, padx=(10,0)) - - def _update_filters(self): - self._create_element_list(search=self.search_var.get(), - type=self.filter_var.get(), - show_disabled=self.show_disabled_var.get()) - self._refresh_show_disabled_text() - - def _reset_filters(self): - self.search_var.set("") - self.filter_var.set("ALL") - self.show_disabled_var.set(True) - self._update_filters() - - def _element_matches_filters(self, element): - # Check enabled status - if not self.filters.get("show_disabled", True): - if hasattr(element, 'enabled') and not element.enabled: - return False - - # Search filter - search = self.filters.get("search", "").lower() - if search: - if not hasattr(element, '_search_cache'): - cache = [] - try: - if getattr(element, 'name', None): - cache.append(element.name.lower()) - p = getattr(element, 'path', None) - if p: - try: - cache.append(os.path.basename(p).lower()) - cache.append(p.lower()) - except (TypeError, AttributeError): - pass - except (AttributeError, TypeError): - pass - element._search_cache = cache - if not any(search in text for text in getattr(element, '_search_cache', [])): - return False - - # Type filter - type_filter = self.filters.get("type", "ALL") - if type_filter != "ALL": - if hasattr(element, 'type') and element.type: - try: - return ConceptType(element.type).value == type_filter - except (ValueError, AttributeError): - return False - return False - - return True + command=self._reset_filters).grid(row=0, column=6, padx=(10, 0)) def _maybe_reposition_toolbar(self, width): if not self._toolbar: @@ -152,6 +86,12 @@ def _maybe_reposition_toolbar(self, width): else: self._toolbar.grid_configure(row=0, column=4, columnspan=2, sticky="ew", padx=10) + def _reset_filters(self): + self.search_var.set("") + self.filter_var.set("ALL") + self.show_disabled_var.set(True) + self._update_filters() + def _refresh_show_disabled_text(self): try: disabled_count = sum(1 for c in getattr(self, 'current_config', []) if getattr(c, 'enabled', True) is False) @@ -165,32 +105,29 @@ def _refresh_show_disabled_text(self): pass -class ConceptWidget(ctk.CTkFrame): - def __init__(self, master, concept, i, open_command, remove_command, clone_command, save_command): - super().__init__( - master=master, width=150, height=170, corner_radius=10, bg_color="transparent" - ) +class CtkConceptWidgetView(BaseConceptWidgetView, ctk.CTkFrame): + + def __init__(self, master, concept, i, open_command, remove_command, clone_command, save_command, controller): + ctk.CTkFrame.__init__(self, master=master, width=150, height=170, corner_radius=10, bg_color="transparent") + BaseConceptWidgetView.__init__(self, ctk_components) self.concept = concept - self.ui_state = UIState(self, concept) - self.image_ui_state = UIState(self, concept.image) - self.text_ui_state = UIState(self, concept.text) + self.ui_state = CtkUIState(self, concept) + self.image_ui_state = CtkUIState(self, concept.image) + self.text_ui_state = CtkUIState(self, concept.text) self.i = i self.grid_rowconfigure(1, weight=1) - # image self.image = ctk.CTkImage( - light_image=self.__get_preview_image(), + light_image=self._get_preview_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) - # name - self.name_label = components.label(self, 1, 0, self.__get_display_name(), pad=5, wraplength=140) + self.name_label = self.components.label(self, 1, 0, self._get_display_name(), pad=5, wraplength=140) - # close button close_button = ctk.CTkButton( master=self, width=20, @@ -202,7 +139,6 @@ def __init__(self, master, concept, i, open_command, remove_command, clone_comma ) close_button.place(x=0, y=0) - # clone button clone_button = ctk.CTkButton( master=self, width=20, @@ -210,11 +146,10 @@ def __init__(self, master, concept, i, open_command, remove_command, clone_comma text="+", corner_radius=2, fg_color="#00C000", - command=lambda: clone_command(self.i, self.__randomize_seed), + command=lambda: clone_command(self.i, controller.randomize_seed), ) clone_button.place(x=25, y=0) - # enabled switch enabled_switch = ctk.CTkSwitch( master=self, width=40, @@ -229,55 +164,10 @@ def __init__(self, master, concept, i, open_command, remove_command, clone_comma lambda event: open_command(self.i, (self.ui_state, self.image_ui_state, self.text_ui_state)) ) - def __randomize_seed(self, concept: ConceptConfig): - concept.seed = ConceptConfig.default_values().seed - return concept - - def __get_display_name(self): - if self.concept.name: - return self.concept.name - elif self.concept.path: - return os.path.basename(self.concept.path) - else: - return "" - def configure_element(self): - self.name_label.configure(text=self.__get_display_name()) - self.image.configure(light_image=self.__get_preview_image()) - try: - if hasattr(self.concept, '_search_cache'): - delattr(self.concept, '_search_cache') - except AttributeError: - pass - - def __get_preview_image(self): - preview_path = "resources/icons/icon.png" - glob_pattern = "**/*.*" if getattr(self.concept, 'include_subdirectories', False) else "*.*" - - concept_path = ConceptWindow.get_concept_path(getattr(self.concept, 'path', None)) - if concept_path: - for path in pathlib.Path(concept_path).glob(glob_pattern): - if any(part.startswith('.') for part in path.relative_to(concept_path).parent.parts): - continue - extension = os.path.splitext(path)[1] - if (path.is_file() - and path_util.is_supported_image_extension(extension) - and not path.name.endswith("-masklabel.png") - and not path.name.endswith("-condlabel.png")): - preview_path = path_util.canonical_join(concept_path, path) - break - try: - image = load_image(preview_path, convert_mode="RGBA") - except (OSError): - image = Image.new("RGBA", (150, 150), (200, 200, 200, 255)) - size = min(image.width, image.height) - image = image.crop(( - (image.width - size) // 2, - (image.height - size) // 2, - (image.width - size) // 2 + size, - (image.height - size) // 2 + size, - )) - return image.resize((150, 150), Image.Resampling.BILINEAR) + self.name_label.configure(text=self._get_display_name()) + self.image.configure(light_image=self._get_preview_image()) + self._clear_search_cache() def place_in_list(self): index = getattr(self, 'visible_index', self.i) diff --git a/modules/ui/CtkConceptWindowView.py b/modules/ui/CtkConceptWindowView.py index f58879d5f..60c0f57fe 100644 --- a/modules/ui/CtkConceptWindowView.py +++ b/modules/ui/CtkConceptWindowView.py @@ -1,94 +1,30 @@ -import fractions -import math -import os -import pathlib -import platform -import random import threading -import time -import traceback - -from modules.util import concept_stats, path_util -from modules.util.config.ConceptConfig import ConceptConfig -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.BalancingStrategy import BalancingStrategy -from modules.util.enum.ConceptType import ConceptType -from modules.util.image_util import load_image -from modules.util.ui import components + +from modules.ui.BaseConceptWindowView import BaseConceptWindowView +from modules.ui.ConceptWindowController import ConceptWindowController +from modules.util.ui import ctk_components from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState - -from mgds.LoadingPipeline import LoadingPipeline -from mgds.OutputPipelineModule import OutputPipelineModule -from mgds.PipelineModule import PipelineModule -from mgds.pipelineModules.CapitalizeTags import CapitalizeTags -from mgds.pipelineModules.DropTags import DropTags -from mgds.pipelineModules.RandomBrightness import RandomBrightness -from mgds.pipelineModules.RandomCircularMaskShrink import ( - RandomCircularMaskShrink, -) -from mgds.pipelineModules.RandomContrast import RandomContrast -from mgds.pipelineModules.RandomFlip import RandomFlip -from mgds.pipelineModules.RandomHue import RandomHue -from mgds.pipelineModules.RandomMaskRotateCrop import RandomMaskRotateCrop -from mgds.pipelineModules.RandomRotate import RandomRotate -from mgds.pipelineModules.RandomSaturation import RandomSaturation -from mgds.pipelineModules.ShuffleTags import ShuffleTags -from mgds.pipelineModuleTypes.RandomAccessPipelineModule import ( - RandomAccessPipelineModule, -) - -import torch -from torchvision.transforms import functional import customtkinter as ctk -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 - - -class InputPipelineModule( - PipelineModule, - RandomAccessPipelineModule, -): - def __init__(self, data: dict): - super().__init__() - self.data = data - - def length(self) -> int: - return 1 - def get_inputs(self) -> list[str]: - return [] - def get_outputs(self) -> list[str]: - return list(self.data.keys()) - - def get_item(self, variation: int, index: int, requested_name: str = None) -> dict: - return self.data - - -class ConceptWindow(ctk.CTkToplevel): +class CtkConceptWindowView(BaseConceptWindowView, ctk.CTkToplevel): def __init__( self, parent, - train_config: TrainConfig, - concept: ConceptConfig, - ui_state: UIState, - image_ui_state: UIState, - text_ui_state: UIState, + controller: ConceptWindowController, + ui_state, + image_ui_state, + text_ui_state, *args, **kwargs, ): - super().__init__(parent, *args, **kwargs) - - self.train_config = train_config + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseConceptWindowView.__init__(self, ctk_components) - self.concept = concept - self.ui_state = ui_state - self.image_ui_state = image_ui_state - self.text_ui_state = text_ui_state + self.controller = controller self.image_preview_file_index = 0 self.preview_augmentations = ctk.BooleanVar(self, True) self.bucket_fig = None @@ -103,197 +39,40 @@ def __init__( tabview = ctk.CTkTabview(self) tabview.grid(row=0, column=0, sticky="nsew") - self.general_tab = self.__general_tab(tabview.add("general"), concept) - self.image_augmentation_tab = self.__image_augmentation_tab(tabview.add("image augmentation")) - self.text_augmentation_tab = self.__text_augmentation_tab(tabview.add("text augmentation")) - self.concept_stats_tab = self.__concept_stats_tab(tabview.add("statistics")) - - #automatic concept scan - self.scan_thread = threading.Thread(target=self.__auto_update_concept_stats, daemon=True) - self.scan_thread.start() - - components.button(self, 1, 0, "ok", self.__ok) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def __general_tab(self, master, concept: ConceptConfig): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, weight=1) - - # name - components.label(frame, 0, 0, "Name", - tooltip="Name of the concept") - components.entry(frame, 0, 1, self.ui_state, "name") - - # enabled - components.label(frame, 1, 0, "Enabled", - tooltip="Enable or disable this concept") - components.switch(frame, 1, 1, self.ui_state, "enabled") - - # concept type - components.label(frame, 2, 0, "Concept Type", - tooltip="STANDARD: Standard finetuning with the sample as training target\n" - "VALIDATION: Use concept for validation instead of training\n" - "PRIOR_PREDICTION: Use the sample to make a prediction using the model as it was before training. This prediction is then used as the training target " - "for the model in training. This can be used as regularisation and to preserve prior model knowledge while finetuning the model on other concepts. " - "Only implemented for LoRA.", - wide_tooltip=True) - components.options(frame, 2, 1, [str(x) for x in list(ConceptType)], self.ui_state, "type") - - # path - components.label(frame, 3, 0, "Path", - tooltip="Path where the training data is located") - components.path_entry(frame, 3, 1, self.ui_state, "path", mode="dir") - 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.path_entry(frame, 4, 2, self.text_ui_state, "prompt_path", mode="file") - - def set_prompt_path_entry_enabled(option: str): - if option == 'concept': - for child in prompt_path_entry.children.values(): - child.configure(state="normal") - else: - for child in prompt_path_entry.children.values(): - child.configure(state="disabled") - - components.options_kv(frame, 4, 1, [ - ("From text file per sample", 'sample'), - ("From single text file", 'concept'), - ("From image file name", 'filename'), - ], self.text_ui_state, "prompt_source", command=set_prompt_path_entry_enabled) - set_prompt_path_entry_enabled(concept.text.prompt_source) - - # include subdirectories - components.label(frame, 5, 0, "Include Subdirectories", - tooltip="Includes images from subdirectories into the dataset") - components.switch(frame, 5, 1, self.ui_state, "include_subdirectories") - - # image variations - components.label(frame, 6, 0, "Image Variations", - tooltip="The number of different image versions to cache if latent caching is enabled.") - components.entry(frame, 6, 1, self.ui_state, "image_variations") - - # text variations - components.label(frame, 7, 0, "Text Variations", - tooltip="The number of different text versions to cache if latent caching is enabled.") - components.entry(frame, 7, 1, self.ui_state, "text_variations") - - # balancing - components.label(frame, 8, 0, "Balancing", - tooltip="The number of samples used during training. Use repeats to multiply the concept, or samples to specify an exact number of samples used in each epoch.") - components.entry(frame, 8, 1, self.ui_state, "balancing") - components.options(frame, 8, 2, [str(x) for x in list(BalancingStrategy)], self.ui_state, "balancing_strategy") - - # loss weight - components.label(frame, 9, 0, "Loss Weight", - tooltip="The loss multiplyer for this concept.") - components.entry(frame, 9, 1, self.ui_state, "loss_weight") - - frame.pack(fill="both", expand=1) - return frame - - def __image_augmentation_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # header - components.label(frame, 0, 1, "Random", - tooltip="Enable this augmentation with random values") - components.label(frame, 0, 2, "Fixed", - tooltip="Enable this augmentation with fixed values") - - # crop jitter - components.label(frame, 1, 0, "Crop Jitter", - tooltip="Enables random cropping of samples") - components.switch(frame, 1, 1, self.image_ui_state, "enable_crop_jitter") - - # random flip - components.label(frame, 2, 0, "Random Flip", - tooltip="Randomly flip the sample during training") - components.switch(frame, 2, 1, self.image_ui_state, "enable_random_flip") - components.switch(frame, 2, 2, self.image_ui_state, "enable_fixed_flip") - - # random rotation - components.label(frame, 3, 0, "Random Rotation", - tooltip="Randomly rotates the sample during training") - components.switch(frame, 3, 1, self.image_ui_state, "enable_random_rotate") - components.switch(frame, 3, 2, self.image_ui_state, "enable_fixed_rotate") - components.entry(frame, 3, 3, self.image_ui_state, "random_rotate_max_angle") - - # random brightness - components.label(frame, 4, 0, "Random Brightness", - tooltip="Randomly adjusts the brightness of the sample during training") - components.switch(frame, 4, 1, self.image_ui_state, "enable_random_brightness") - components.switch(frame, 4, 2, self.image_ui_state, "enable_fixed_brightness") - components.entry(frame, 4, 3, self.image_ui_state, "random_brightness_max_strength") - - # random contrast - components.label(frame, 5, 0, "Random Contrast", - tooltip="Randomly adjusts the contrast of the sample during training") - components.switch(frame, 5, 1, self.image_ui_state, "enable_random_contrast") - components.switch(frame, 5, 2, self.image_ui_state, "enable_fixed_contrast") - components.entry(frame, 5, 3, self.image_ui_state, "random_contrast_max_strength") - - # random saturation - components.label(frame, 6, 0, "Random Saturation", - tooltip="Randomly adjusts the saturation of the sample during training") - components.switch(frame, 6, 1, self.image_ui_state, "enable_random_saturation") - components.switch(frame, 6, 2, self.image_ui_state, "enable_fixed_saturation") - components.entry(frame, 6, 3, self.image_ui_state, "random_saturation_max_strength") - - # random hue - components.label(frame, 7, 0, "Random Hue", - tooltip="Randomly adjusts the hue of the sample during training") - components.switch(frame, 7, 1, self.image_ui_state, "enable_random_hue") - components.switch(frame, 7, 2, self.image_ui_state, "enable_fixed_hue") - components.entry(frame, 7, 3, self.image_ui_state, "random_hue_max_strength") - - # random circular mask shrink - components.label(frame, 8, 0, "Circular Mask Generation", - tooltip="Automatically create circular masks for masked training") - components.switch(frame, 8, 1, self.image_ui_state, "enable_random_circular_mask_shrink") - - # random rotate and crop - components.label(frame, 9, 0, "Random Rotate and Crop", - tooltip="Randomly rotate the training samples and crop to the masked region") - components.switch(frame, 9, 1, self.image_ui_state, "enable_random_mask_rotate_crop") - - # circular mask generation - components.label(frame, 10, 0, "Resolution Override", - tooltip="Override the resolution for this concept. Optionally specify multiple resolutions separated by a comma, or a single exact resolution in the format x") - components.switch(frame, 10, 2, self.image_ui_state, "enable_resolution_override") - components.entry(frame, 10, 3, self.image_ui_state, "resolution_override") + # general tab + general_frame = ctk.CTkScrollableFrame(tabview.add("general"), fg_color="transparent") + general_frame.grid_columnconfigure(1, weight=1) + general_frame.grid_columnconfigure(2, weight=1) + self.build_general_tab(general_frame, controller, ui_state, text_ui_state) + general_frame.pack(fill="both", expand=1) + + # image augmentation tab + image_aug_master = tabview.add("image augmentation") + image_aug_frame = ctk.CTkScrollableFrame(image_aug_master, fg_color="transparent") + image_aug_frame.grid_columnconfigure(0, weight=0) + image_aug_frame.grid_columnconfigure(1, weight=0) + image_aug_frame.grid_columnconfigure(2, weight=0) + image_aug_frame.grid_columnconfigure(3, weight=1) + self.build_image_augmentation_tab(image_aug_frame, controller, image_ui_state) # image - image_preview, filename_preview, caption_preview = self.__get_preview_image() + image_preview, filename_preview, caption_preview = controller.get_preview_image(self.image_preview_file_index, self.preview_augmentations.get()) self.image = ctk.CTkImage( light_image=image_preview, size=image_preview.size, ) - image_label = ctk.CTkLabel(master=frame, text="", image=self.image, height=300, width=300) + image_label = ctk.CTkLabel(master=image_aug_frame, text="", image=self.image, height=300, width=300) image_label.grid(row=0, column=4, rowspan=6) # refresh preview - update_button_frame = ctk.CTkFrame(master=frame, corner_radius=0, fg_color="transparent") + update_button_frame = ctk.CTkFrame(master=image_aug_frame, corner_radius=0, fg_color="transparent") update_button_frame.grid(row=6, column=4, rowspan=6, sticky="nsew") update_button_frame.grid_columnconfigure(1, weight=1) - prev_preview_button = components.button(update_button_frame, 0, 0, "<", command=self.__prev_image_preview) - components.button(update_button_frame, 0, 1, "Update Preview", command=self.__update_image_preview) - next_preview_button = components.button(update_button_frame, 0, 2, ">", command=self.__next_image_preview) - preview_augmentations_switch = ctk.CTkSwitch(update_button_frame, text="Show Augmentations", variable=self.preview_augmentations, command=self.__update_image_preview) + prev_preview_button = self.components.button(update_button_frame, 0, 0, "<", command=self._prev_image_preview) + self.components.button(update_button_frame, 0, 1, "Update Preview", command=self._update_image_preview) + next_preview_button = self.components.button(update_button_frame, 0, 2, ">", command=self._next_image_preview) + preview_augmentations_switch = ctk.CTkSwitch(update_button_frame, text="Show Augmentations", variable=self.preview_augmentations, command=self._update_image_preview) preview_augmentations_switch.grid(row=1, column=0, columnspan=3, padx=5, pady=5) prev_preview_button.configure(width=40) @@ -307,205 +86,24 @@ def __image_augmentation_tab(self, master): self.caption_preview.configure(state="disabled") self.caption_preview.grid(row=3, column=0, columnspan=3, rowspan=3) - frame.pack(fill="both", expand=1) - return frame - - def __text_augmentation_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # tag shuffling - components.label(frame, 0, 0, "Tag Shuffling", - tooltip="Enables tag shuffling") - components.switch(frame, 0, 1, self.text_ui_state, "enable_tag_shuffling") - - # keep tag count - components.label(frame, 1, 0, "Tag Delimiter", - tooltip="The delimiter between tags") - components.entry(frame, 1, 1, self.text_ui_state, "tag_delimiter") - - # keep tag count - components.label(frame, 2, 0, "Keep Tag Count", - tooltip="The number of tags at the start of the caption that are not shuffled or dropped") - components.entry(frame, 2, 1, self.text_ui_state, "keep_tags_count") - - # tag dropout - components.label(frame, 3, 0, "Tag Dropout", - tooltip="Enables random dropout for tags in the captions.") - components.switch(frame, 3, 1, self.text_ui_state, "tag_dropout_enable") - components.label(frame, 4, 0, "Dropout Mode", - tooltip="Method used to drop captions. 'Full' will drop the entire caption past the 'kept' tags with a certain probability, 'Random' will drop individual tags with the set probability, and 'Random Weighted' will linearly increase the probability of dropping tags, more likely to preseve tags near the front with full probability to drop at the end.") - components.options_kv(frame, 4, 1, [ - ("Full", 'FULL'), - ("Random", 'RANDOM'), - ("Random Weighted", 'RANDOM WEIGHTED'), - ], self.text_ui_state, "tag_dropout_mode", None) - components.label(frame, 4, 2, "Probability", - tooltip="Probability to drop tags, from 0 to 1.") - components.entry(frame, 4, 3, self.text_ui_state, "tag_dropout_probability") - - components.label(frame, 5, 0, "Special Dropout Tags", - tooltip="List of tags which will be whitelisted/blacklisted by dropout. 'Whitelist' tags will never be dropped but all others may be, 'Blacklist' tags may be dropped but all others will never be, 'None' may drop any tags. Can specify either a delimiter-separated list in the field, or a file path to a .txt or .csv file with entries separated by newlines.") - components.options_kv(frame, 5, 1, [ - ("None", 'NONE'), - ("Blacklist", 'BLACKLIST'), - ("Whitelist", 'WHITELIST'), - ], self.text_ui_state, "tag_dropout_special_tags_mode", None) - components.entry(frame, 5, 2, self.text_ui_state, "tag_dropout_special_tags") - components.label(frame, 6, 0, "Special Tags Regex", - tooltip="Interpret special tags with regex, such as 'photo.*' to match 'photo, photograph, photon' but not 'telephoto'. Includes exception for '/(' and '/)' syntax found in many booru/e6 tags.") - components.switch(frame, 6, 1, self.text_ui_state, "tag_dropout_special_tags_regex") - - #capitalization randomization - components.label(frame, 7, 0, "Randomize Capitalization", - tooltip="Enables randomization of capitalization for tags in the caption.") - components.switch(frame, 7, 1, self.text_ui_state, "caps_randomize_enable") - components.label(frame, 7, 2, "Force Lowercase", - tooltip="If enabled, converts the caption to lowercase before any further processing.") - components.switch(frame, 7, 3, self.text_ui_state, "caps_randomize_lowercase") - - components.label(frame, 8, 0, "Captialization Mode", - tooltip="Comma-separated list of types of capitalization randomization to perform. 'capslock' for ALL CAPS, 'title' for First Letter Of Every Word, 'first' for First word only, 'random' for rAndOMiZeD lEtTERs.") - components.entry(frame, 8, 1, self.text_ui_state, "caps_randomize_mode") - components.label(frame, 8, 2, "Probability", - tooltip="Probability to randomize capitialization of each tag, from 0 to 1.") - components.entry(frame, 8, 3, self.text_ui_state, "caps_randomize_probability") - - frame.pack(fill="both", expand=1) - return frame - - def __concept_stats_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=150) - frame.grid_columnconfigure(1, weight=0, minsize=150) - frame.grid_columnconfigure(2, weight=0, minsize=150) - frame.grid_columnconfigure(3, weight=0, minsize=150) - - self.cancel_scan_flag = threading.Event() - - #file size - self.file_size_label = components.label(frame, 1, 0, "Total Size", pad=0, - tooltip="Total size of all image, mask, and caption files in MB") - self.file_size_label.configure(font=ctk.CTkFont(underline=True)) - self.file_size_preview = components.label(frame, 2, 0, pad=0, text="-") - - #subdirectory count - self.dir_count_label = components.label(frame, 1, 1, "Directories", pad=0, - tooltip="Total number of directories including and under (if 'include subdirectories' is enabled) the main concept directory") - self.dir_count_label.configure(font=ctk.CTkFont(underline=True)) - self.dir_count_preview = components.label(frame, 2, 1, pad=0, text="-") - - #basic img/vid stats - count of each type in the concept - #the \n at the start of the label gives it better vertical spacing with other rows - self.image_count_label = components.label(frame, 3, 0, "\nTotal Images", pad=0, - tooltip="Total number of image files, any of the extensions " + str(path_util.SUPPORTED_IMAGE_EXTENSIONS) + ", excluding '-masklabel.png and -condlabel.png'") - self.image_count_label.configure(font=ctk.CTkFont(underline=True)) - self.image_count_preview = components.label(frame, 4, 0, pad=0, text="-") - self.video_count_label = components.label(frame, 3, 1, "\nTotal Videos", pad=0, - tooltip="Total number of video files, any of the extensions " + str(path_util.SUPPORTED_VIDEO_EXTENSIONS)) - self.video_count_label.configure(font=ctk.CTkFont(underline=True)) - self.video_count_preview = components.label(frame, 4, 1, pad=0, text="-") - self.mask_count_label = components.label(frame, 3, 2, "\nTotal Masks", pad=0, - tooltip="Total number of mask files, any file ending in '-masklabel.png'") - self.mask_count_label.configure(font=ctk.CTkFont(underline=True)) - self.mask_count_preview = components.label(frame, 4, 2, pad=0, text="-") - self.caption_count_label = components.label(frame, 3, 3, "\nTotal Captions", pad=0, - tooltip="Total number of caption files, any .txt file. With advanced scan, includes the total number of captions on separate lines across all files in parentheses.") - self.caption_count_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_count_preview = components.label(frame, 4, 3, pad=0, text="-") - - #advanced img/vid stats - how many img/vid files have a mask or caption of the same name - self.image_count_mask_label = components.label(frame, 5, 0, "\nImages with Masks", pad=0, - tooltip="Total number of image files with an associated mask") - self.image_count_mask_label.configure(font=ctk.CTkFont(underline=True)) - self.image_count_mask_preview = components.label(frame, 6, 0, pad=0, text="-") - self.mask_count_label_unpaired = components.label(frame, 5, 1, "\nUnpaired Masks", pad=0, - tooltip="Total number of mask files which lack a corresponding image file - if >0, check your data set!") - self.mask_count_label_unpaired.configure(font=ctk.CTkFont(underline=True)) - self.mask_count_preview_unpaired = components.label(frame, 6, 1, pad=0, text="-") - #currently no masks for videos? - - self.image_count_caption_label = components.label(frame, 7, 0, "\nImages with Captions", pad=0, - tooltip="Total number of image files with an associated caption") - self.image_count_caption_label.configure(font=ctk.CTkFont(underline=True)) - self.image_count_caption_preview = components.label(frame, 8, 0, pad=0, text="-") - self.video_count_caption_label = components.label(frame, 7, 1, "\nVideos with Captions", pad=0, - tooltip="Total number of video files with an associated caption") - self.video_count_caption_label.configure(font=ctk.CTkFont(underline=True)) - self.video_count_caption_preview = components.label(frame, 8, 1, pad=0, text="-") - self.caption_count_label_unpaired = components.label(frame, 7, 2, "\nUnpaired Captions", pad=0, - tooltip="Total number of caption files which lack a corresponding image file - if >0, check your data set! If using 'from file name' or 'from single text file' then this can be ignored.") - self.caption_count_label_unpaired.configure(font=ctk.CTkFont(underline=True)) - self.caption_count_preview_unpaired = components.label(frame, 8, 2, pad=0, text="-") - - #resolution info - self.pixel_max_label = components.label(frame, 9, 0, "\nMax Pixels", pad=0, - tooltip="Largest image in the concept by number of pixels (width * height)") - self.pixel_max_label.configure(font=ctk.CTkFont(underline=True)) - self.pixel_max_preview = components.label(frame, 10, 0, pad=0, text="-", wraplength=150) - self.pixel_avg_label = components.label(frame, 9, 1, "\nAvg Pixels", pad=0, - tooltip="Average size of images in the concept by number of pixels (width * height)") - self.pixel_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.pixel_avg_preview = components.label(frame, 10, 1, pad=0, text="-", wraplength=150) - self.pixel_min_label = components.label(frame, 9, 2, "\nMin Pixels", pad=0, - tooltip="Smallest image in the concept by number of pixels (width * height)") - self.pixel_min_label.configure(font=ctk.CTkFont(underline=True)) - self.pixel_min_preview = components.label(frame, 10, 2, pad=0, text="-", wraplength=150) - - #video length info - self.length_max_label = components.label(frame, 11, 0, "\nMax Length", pad=0, - tooltip="Longest video in the concept by number of frames") - self.length_max_label.configure(font=ctk.CTkFont(underline=True)) - self.length_max_preview = components.label(frame, 12, 0, pad=0, text="-", wraplength=150) - self.length_avg_label = components.label(frame, 11, 1, "\nAvg Length", pad=0, - tooltip="Average length of videos in the concept by number of frames") - self.length_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.length_avg_preview = components.label(frame, 12, 1, pad=0, text="-", wraplength=150) - self.length_min_label = components.label(frame, 11, 2, "\nMin Length", pad=0, - tooltip="Shortest video in the concept by number of frames") - self.length_min_label.configure(font=ctk.CTkFont(underline=True)) - self.length_min_preview = components.label(frame, 12, 2, pad=0, text="-", wraplength=150) - - #video fps info - self.fps_max_label = components.label(frame, 13, 0, "\nMax FPS", pad=0, - tooltip="Video in concept with highest fps") - self.fps_max_label.configure(font=ctk.CTkFont(underline=True)) - self.fps_max_preview = components.label(frame, 14, 0, pad=0, text="-", wraplength=150) - self.fps_avg_label = components.label(frame, 13, 1, "\nAvg FPS", pad=0, - tooltip="Average fps of videos in the concept") - self.fps_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.fps_avg_preview = components.label(frame, 14, 1, pad=0, text="-", wraplength=150) - self.fps_min_label = components.label(frame, 13, 2, "\nMin FPS", pad=0, - tooltip="Video in concept with the lowest fps") - self.fps_min_label.configure(font=ctk.CTkFont(underline=True)) - self.fps_min_preview = components.label(frame, 14, 2, pad=0, text="-", wraplength=150) - - #caption info - self.caption_max_label = components.label(frame, 15, 0, "\nMax Caption Length", pad=0, - tooltip="Largest caption in concept by character count. For token count, assume ~2 tokens/word") - self.caption_max_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_max_preview = components.label(frame, 16, 0, pad=0, text="-", wraplength=150) - self.caption_avg_label = components.label(frame, 15, 1, "\nAvg Caption Length", pad=0, - tooltip="Average length of caption in concept by character count. For token count, assume ~2 tokens/word") - self.caption_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_avg_preview = components.label(frame, 16, 1, pad=0, text="-", wraplength=150) - self.caption_min_label = components.label(frame, 15, 2, "\nMin Caption Length", pad=0, - tooltip="Smallest caption in concept by character count. For token count, assume ~2 tokens/word") - self.caption_min_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_min_preview = components.label(frame, 16, 2, pad=0, text="-", wraplength=150) - - #aspect bucket info - self.aspect_bucket_label = components.label(frame, 17, 0, "\nAspect Bucketing", pad=0, - tooltip="Graph of all possible buckets and the number of images in each one, defined as height/width. Buckets range from 0.25 (4:1 extremely wide) to 4 (1:4 extremely tall). \ - Images which don't match a bucket exactly are cropped to the nearest one.") - self.aspect_bucket_label.configure(font=ctk.CTkFont(underline=True)) - self.small_bucket_label = components.label(frame, 17, 1, "\nSmallest Buckets", pad=0, - tooltip="Image buckets with the least nonzero total images - if 'batch size' is larger than this, these images will be ignored during training! See the wiki for more details.") - self.small_bucket_label.configure(font=ctk.CTkFont(underline=True)) - self.small_bucket_preview = components.label(frame, 18, 1, pad=0, text="-") + image_aug_frame.pack(fill="both", expand=1) + + # text augmentation tab + text_aug_frame = ctk.CTkScrollableFrame(tabview.add("text augmentation"), fg_color="transparent") + text_aug_frame.grid_columnconfigure(0, weight=0) + text_aug_frame.grid_columnconfigure(1, weight=0) + text_aug_frame.grid_columnconfigure(2, weight=0) + text_aug_frame.grid_columnconfigure(3, weight=1) + self.build_text_augmentation_tab(text_aug_frame, controller, text_ui_state) + text_aug_frame.pack(fill="both", expand=1) + + # statistics tab + stats_frame = ctk.CTkScrollableFrame(tabview.add("statistics"), fg_color="transparent") + stats_frame.grid_columnconfigure(0, weight=0, minsize=150) + stats_frame.grid_columnconfigure(1, weight=0, minsize=150) + stats_frame.grid_columnconfigure(2, weight=0, minsize=150) + stats_frame.grid_columnconfigure(3, weight=0, minsize=150) + self.build_concept_stats_tab(stats_frame, controller) #aspect bucketing plot, mostly copied from timestep preview graph appearance_mode = AppearanceModeTracker.get_mode() @@ -518,7 +116,7 @@ def __concept_stats_tab(self, master): assert self.bucket_fig is None self.bucket_fig, self.bucket_ax = plt.subplots(figsize=(7,3)) - self.canvas = FigureCanvasTkAgg(self.bucket_fig, master=frame) + self.canvas = FigureCanvasTkAgg(self.bucket_fig, master=stats_frame) self.canvas.get_tk_widget().grid(row=19, column=0, columnspan=4, rowspan=2) self.bucket_fig.tight_layout() self.bucket_fig.subplots_adjust(bottom=0.15) @@ -534,28 +132,29 @@ 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) - #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") - self.refresh_advanced_stats_button = components.button(master=frame, row=0, column=1, text="Refresh Advanced", command=lambda: self.__get_concept_stats_threaded(True, 9999), - tooltip="Reload advanced statistics for the concept directory") #run "basic" scan first before "advanced", seems to help the system cache the directories and run faster - self.cancel_stats_button = components.button(master=frame, row=0, column=2, text="Abort Scan", command=lambda: self.__cancel_concept_stats(), - tooltip="Stop the currently running scan if it's taking a long time - advanced scan will be slow on large folders and on HDDs") - self.processing_time = components.label(frame, 0, 3, text="-", tooltip="Time taken to process concept directory") + stats_frame.pack(fill="both", expand=1) - frame.pack(fill="both", expand=1) - return frame + #automatic concept scan + self.scan_thread = threading.Thread(target=controller.auto_update_concept_stats, args=[self], daemon=True) + self.scan_thread.start() - def __prev_image_preview(self): + self.components.button(self, 1, 0, "ok", self._ok) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + def _prev_image_preview(self): self.image_preview_file_index = max(self.image_preview_file_index - 1, 0) - self.__update_image_preview() + self._update_image_preview() - def __next_image_preview(self): + def _next_image_preview(self): self.image_preview_file_index += 1 - self.__update_image_preview() + self._update_image_preview() - def __update_image_preview(self): - image_preview, filename_preview, caption_preview = self.__get_preview_image() + def _update_image_preview(self): + image_preview, filename_preview, caption_preview = self.controller.get_preview_image(self.image_preview_file_index, self.preview_augmentations.get()) self.image.configure(light_image=image_preview, size=image_preview.size) self.filename_preview.configure(text=filename_preview) self.caption_preview.configure(state="normal") @@ -563,366 +162,6 @@ def __update_image_preview(self): self.caption_preview.insert(index="1.0", text=caption_preview) self.caption_preview.configure(state="disabled") - @staticmethod - def get_concept_path(path: str) -> str | None: - if os.path.isdir(path): - return path - try: - #don't download, only check if available locally: - return huggingface_hub.snapshot_download(repo_id=path, repo_type="dataset", local_files_only=True) - except Exception: - return None - - def __download_dataset(self): - try: - huggingface_hub.login(token=self.train_config.secrets.huggingface_token, new_session=False) - huggingface_hub.snapshot_download(repo_id=self.concept.path, repo_type="dataset") - except Exception: - traceback.print_exc() - - def __download_dataset_threaded(self): - download_thread = threading.Thread(target=self.__download_dataset, daemon=True) - download_thread.start() - - def _read_text_file_for_preview(self, file_path: str) -> str: - empty_msg = "[Empty prompt]" - try: - with open(file_path, "r") as f: - if self.preview_augmentations.get(): - lines = [line.strip() for line in f if line.strip()] - return random.choice(lines) if lines else empty_msg - content = f.read().strip() - return content if content else empty_msg - except FileNotFoundError: - return "File not found, please check the path" - except IsADirectoryError: - return "[Provided path is a directory, please correct the caption path]" - except PermissionError: - if platform.system() == "Windows": - return "[Permission denied, please check the file permissions or Windows Defender settings]" - else: - return "[Permission denied, please check the file permissions]" - except UnicodeDecodeError: - return "[Invalid file encoding. This should not happen, please report this issue]" - - def __get_preview_image(self): - preview_image_path = "resources/icons/icon.png" - file_index = -1 - glob_pattern = "**/*.*" if self.concept.include_subdirectories else "*.*" - - concept_path = self.get_concept_path(self.concept.path) - if concept_path: - for path in pathlib.Path(concept_path).glob(glob_pattern): - if any(part.startswith('.') for part in path.relative_to(concept_path).parent.parts): - continue - extension = os.path.splitext(path)[1] - if path.is_file() and path_util.is_supported_image_extension(extension) \ - and not path.name.endswith("-masklabel.png") and not path.name.endswith("-condlabel.png"): - preview_image_path = path_util.canonical_join(concept_path, path) - file_index += 1 - if file_index == self.image_preview_file_index: - break - - image = load_image(preview_image_path, 'RGB') - image_tensor = functional.to_tensor(image) - - splitext = os.path.splitext(preview_image_path) - preview_mask_path = path_util.canonical_join(splitext[0] + "-masklabel.png") - if not os.path.isfile(preview_mask_path): - preview_mask_path = None - - if preview_mask_path: - mask = Image.open(preview_mask_path).convert("L") - mask_tensor = functional.to_tensor(mask) - else: - mask_tensor = torch.ones((1, image_tensor.shape[1], image_tensor.shape[2])) - - source = self.concept.text.prompt_source - preview_p = pathlib.Path(preview_image_path) - if source == "filename": - prompt_output = preview_p.stem or "[Empty prompt]" - else: - file_map = { - "sample": preview_p.with_suffix(".txt"), - "concept": pathlib.Path(self.concept.text.prompt_path) if self.concept.text.prompt_path else None, - } - file_path = file_map.get(source) - prompt_output = self._read_text_file_for_preview(str(file_path)) if file_path else "[Empty prompt]" - - modules = [] - if self.preview_augmentations.get(): - input_module = InputPipelineModule({ - 'true': True, - 'image': image_tensor, - 'mask': mask_tensor, - 'enable_random_flip': self.concept.image.enable_random_flip, - 'enable_fixed_flip': self.concept.image.enable_fixed_flip, - 'enable_random_rotate': self.concept.image.enable_random_rotate, - 'enable_fixed_rotate': self.concept.image.enable_fixed_rotate, - 'random_rotate_max_angle': self.concept.image.random_rotate_max_angle, - 'enable_random_brightness': self.concept.image.enable_random_brightness, - 'enable_fixed_brightness': self.concept.image.enable_fixed_brightness, - 'random_brightness_max_strength': self.concept.image.random_brightness_max_strength, - 'enable_random_contrast': self.concept.image.enable_random_contrast, - 'enable_fixed_contrast': self.concept.image.enable_fixed_contrast, - 'random_contrast_max_strength': self.concept.image.random_contrast_max_strength, - 'enable_random_saturation': self.concept.image.enable_random_saturation, - 'enable_fixed_saturation': self.concept.image.enable_fixed_saturation, - 'random_saturation_max_strength': self.concept.image.random_saturation_max_strength, - 'enable_random_hue': self.concept.image.enable_random_hue, - 'enable_fixed_hue': self.concept.image.enable_fixed_hue, - 'random_hue_max_strength': self.concept.image.random_hue_max_strength, - 'enable_random_circular_mask_shrink': self.concept.image.enable_random_circular_mask_shrink, - 'enable_random_mask_rotate_crop': self.concept.image.enable_random_mask_rotate_crop, - - 'prompt' : prompt_output, - 'tag_dropout_enable' : self.concept.text.tag_dropout_enable, - 'tag_dropout_probability' : self.concept.text.tag_dropout_probability, - 'tag_dropout_mode' : self.concept.text.tag_dropout_mode, - 'tag_dropout_special_tags' : self.concept.text.tag_dropout_special_tags, - 'tag_dropout_special_tags_mode' : self.concept.text.tag_dropout_special_tags_mode, - 'tag_delimiter' : self.concept.text.tag_delimiter, - 'keep_tags_count' : self.concept.text.keep_tags_count, - 'tag_dropout_special_tags_regex' : self.concept.text.tag_dropout_special_tags_regex, - 'caps_randomize_enable' : self.concept.text.caps_randomize_enable, - 'caps_randomize_probability' : self.concept.text.caps_randomize_probability, - 'caps_randomize_mode' : self.concept.text.caps_randomize_mode, - 'caps_randomize_lowercase' : self.concept.text.caps_randomize_lowercase, - 'enable_tag_shuffling' : self.concept.text.enable_tag_shuffling, - }) - - circular_mask_shrink = RandomCircularMaskShrink(mask_name='mask', shrink_probability=1.0, shrink_factor_min=0.2, shrink_factor_max=1.0, enabled_in_name='enable_random_circular_mask_shrink') - random_mask_rotate_crop = RandomMaskRotateCrop(mask_name='mask', additional_names=['image'], min_size=512, min_padding_percent=10, max_padding_percent=30, max_rotate_angle=20, enabled_in_name='enable_random_mask_rotate_crop') - random_flip = RandomFlip(names=['image', 'mask'], enabled_in_name='enable_random_flip', fixed_enabled_in_name='enable_fixed_flip') - random_rotate = RandomRotate(names=['image', 'mask'], enabled_in_name='enable_random_rotate', fixed_enabled_in_name='enable_fixed_rotate', max_angle_in_name='random_rotate_max_angle') - random_brightness = RandomBrightness(names=['image'], enabled_in_name='enable_random_brightness', fixed_enabled_in_name='enable_fixed_brightness', max_strength_in_name='random_brightness_max_strength') - random_contrast = RandomContrast(names=['image'], enabled_in_name='enable_random_contrast', fixed_enabled_in_name='enable_fixed_contrast', max_strength_in_name='random_contrast_max_strength') - random_saturation = RandomSaturation(names=['image'], enabled_in_name='enable_random_saturation', fixed_enabled_in_name='enable_fixed_saturation', max_strength_in_name='random_saturation_max_strength') - random_hue = RandomHue(names=['image'], enabled_in_name='enable_random_hue', fixed_enabled_in_name='enable_fixed_hue', max_strength_in_name='random_hue_max_strength') - drop_tags = DropTags(text_in_name='prompt', enabled_in_name='tag_dropout_enable', probability_in_name='tag_dropout_probability', dropout_mode_in_name='tag_dropout_mode', - special_tags_in_name='tag_dropout_special_tags', special_tag_mode_in_name='tag_dropout_special_tags_mode', delimiter_in_name='tag_delimiter', - keep_tags_count_in_name='keep_tags_count', text_out_name='prompt', regex_enabled_in_name='tag_dropout_special_tags_regex') - caps_randomize = CapitalizeTags(text_in_name='prompt', enabled_in_name='caps_randomize_enable', probability_in_name='caps_randomize_probability', - capitalize_mode_in_name='caps_randomize_mode', delimiter_in_name='tag_delimiter', convert_lowercase_in_name='caps_randomize_lowercase', text_out_name='prompt') - shuffle_tags = ShuffleTags(text_in_name='prompt', enabled_in_name='enable_tag_shuffling', delimiter_in_name='tag_delimiter', keep_tags_count_in_name='keep_tags_count', text_out_name='prompt') - output_module = OutputPipelineModule(['image', 'mask', 'prompt']) - - modules = [ - input_module, - circular_mask_shrink, - random_mask_rotate_crop, - random_flip, - random_rotate, - random_brightness, - random_contrast, - random_saturation, - random_hue, - drop_tags, - caps_randomize, - shuffle_tags, - output_module, - ] - - pipeline = LoadingPipeline( - device=torch.device('cpu'), - modules=modules, - batch_size=1, - seed=random.randint(0, 2**30), - state=None, - initial_epoch=0, - initial_index=0, - ) - - data = pipeline.__next__() - image_tensor = data['image'] - mask_tensor = data['mask'] - prompt_output = data['prompt'] - - filename_output = os.path.basename(preview_image_path) - - mask_tensor = torch.clamp(mask_tensor, 0.3, 1) - image_tensor = image_tensor * mask_tensor - - image = functional.to_pil_image(image_tensor) - - image.thumbnail((300, 300)) - - return image, filename_output, prompt_output - - def __update_concept_stats(self): - #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") - - #directory count - self.dir_count_preview.configure(text=self.concept.concept_stats["directory_count"]) - - #image count - self.image_count_preview.configure(text=self.concept.concept_stats["image_count"]) - self.image_count_mask_preview.configure(text=self.concept.concept_stats["image_with_mask_count"]) - self.image_count_caption_preview.configure(text=self.concept.concept_stats["image_with_caption_count"]) - - #video count - self.video_count_preview.configure(text=self.concept.concept_stats["video_count"]) - #self.video_count_mask_preview.configure(text=self.concept.concept_stats["video_with_mask_count"]) - self.video_count_caption_preview.configure(text=self.concept.concept_stats["video_with_caption_count"]) - - #mask count - self.mask_count_preview.configure(text=self.concept.concept_stats["mask_count"]) - self.mask_count_preview_unpaired.configure(text=self.concept.concept_stats["unpaired_masks"]) - - #caption count - if self.concept.concept_stats["subcaption_count"] > 0: - self.caption_count_preview.configure(text=f'{self.concept.concept_stats["caption_count"]} ({self.concept.concept_stats["subcaption_count"]})') - else: - self.caption_count_preview.configure(text=self.concept.concept_stats["caption_count"]) - self.caption_count_preview_unpaired.configure(text=self.concept.concept_stats["unpaired_captions"]) - - #resolution info - max_pixels = self.concept.concept_stats["max_pixels"] - avg_pixels = self.concept.concept_stats["avg_pixels"] - min_pixels = self.concept.concept_stats["min_pixels"] - - if any(isinstance(x, str) for x in [max_pixels, avg_pixels, min_pixels]) or self.concept.concept_stats["image_count"] == 0: #will be str if adv stats were not taken - self.pixel_max_preview.configure(text="-") - self.pixel_avg_preview.configure(text="-") - self.pixel_min_preview.configure(text="-") - else: - #formatted as (#pixels/1000000) MP, width x height, \n filename - self.pixel_max_preview.configure(text=f'{str(round(max_pixels[0]/1000000, 2))} MP, {max_pixels[2]}\n{max_pixels[1]}') - self.pixel_avg_preview.configure(text=f'{str(round(avg_pixels/1000000, 2))} MP, ~{int(math.sqrt(avg_pixels))}w x {int(math.sqrt(avg_pixels))}h') - self.pixel_min_preview.configure(text=f'{str(round(min_pixels[0]/1000000, 2))} MP, {min_pixels[2]}\n{min_pixels[1]}') - - #video length and fps info - max_length = self.concept.concept_stats["max_length"] - avg_length = self.concept.concept_stats["avg_length"] - min_length = self.concept.concept_stats["min_length"] - max_fps = self.concept.concept_stats["max_fps"] - avg_fps = self.concept.concept_stats["avg_fps"] - min_fps = self.concept.concept_stats["min_fps"] - - if any(isinstance(x, str) for x in [max_length, avg_length, min_length]) or self.concept.concept_stats["video_count"] == 0: #will be str if adv stats were not taken - self.length_max_preview.configure(text="-") - self.length_avg_preview.configure(text="-") - self.length_min_preview.configure(text="-") - self.fps_max_preview.configure(text="-") - self.fps_avg_preview.configure(text="-") - self.fps_min_preview.configure(text="-") - else: - #formatted as (#frames) frames \n filename - self.length_max_preview.configure(text=f'{int(max_length[0])} frames\n{max_length[1]}') - self.length_avg_preview.configure(text=f'{int(avg_length)} frames') - self.length_min_preview.configure(text=f'{int(min_length[0])} frames\n{min_length[1]}') - #formatted as (#fps) fps \n filename - self.fps_max_preview.configure(text=f'{int(max_fps[0])} fps\n{max_fps[1]}') - self.fps_avg_preview.configure(text=f'{int(avg_fps)} fps') - self.fps_min_preview.configure(text=f'{int(min_fps[0])} fps\n{min_fps[1]}') - - #caption info - max_caption_length = self.concept.concept_stats["max_caption_length"] - avg_caption_length = self.concept.concept_stats["avg_caption_length"] - min_caption_length = self.concept.concept_stats["min_caption_length"] - - if any(isinstance(x, str) for x in [max_caption_length, avg_caption_length, min_caption_length]) or self.concept.concept_stats["caption_count"] == 0: #will be str if adv stats were not taken - self.caption_max_preview.configure(text="-") - self.caption_avg_preview.configure(text="-") - self.caption_min_preview.configure(text="-") - else: - #formatted as (#chars) chars, (#words) words, \n filename - self.caption_max_preview.configure(text=f'{max_caption_length[0]} chars, {max_caption_length[2]} words\n{max_caption_length[1]}') - self.caption_avg_preview.configure(text=f'{int(avg_caption_length[0])} chars, {int(avg_caption_length[1])} words') - self.caption_min_preview.configure(text=f'{min_caption_length[0]} chars, {min_caption_length[2]} words\n{min_caption_length[1]}') - - #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) - - self.bucket_ax.cla() - aspects = [str(x) for x in list(aspect_buckets.keys())] - aspect_ratios = [self.decimal_to_aspect_ratio(x) for x in list(aspect_buckets.keys())] - counts = list(aspect_buckets.values()) - b = self.bucket_ax.bar(aspect_ratios, counts) - self.bucket_ax.bar_label(b, color=self.text_color) - sec = self.bucket_ax.secondary_xaxis(location=-0.1) - sec.spines["bottom"].set_linewidth(0) - sec.set_xticks([0, (len(aspects)-1)/2, len(aspects)-1], labels=["Wide", "Square", "Tall"]) - sec.tick_params('x', length=0) - self.canvas.draw() - - def decimal_to_aspect_ratio(self, value : float): - #find closest fraction to decimal aspect value and convert to a:b format - aspect_fraction = fractions.Fraction(value).limit_denominator(16) - aspect_string = f'{aspect_fraction.denominator}:{aspect_fraction.numerator}' - return aspect_string - - def __get_concept_stats(self, advanced_checks: bool, wait_time: float): - start_time = time.perf_counter() - last_update = time.perf_counter() - self.cancel_scan_flag.clear() - self.concept_stats_tab.after(0, self.__disable_scan_buttons) - concept_path = self.get_concept_path(self.concept.path) - - if not concept_path: - print(f"Unable to get statistics for concept path: {self.concept.path}") - self.concept_stats_tab.after(0, self.__enable_scan_buttons) - return - subfolders = [concept_path] - - 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: - 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 - subfolders.extend([f for f in os.scandir(path) if f.is_dir() and not f.name.startswith('.')]) - self.concept.concept_stats = stats_dict - #update GUI approx every half second - if time.perf_counter() > (last_update + 0.5): - 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) - - 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): - self.refresh_basic_stats_button.configure(state="disabled") - self.refresh_advanced_stats_button.configure(state="disabled") - - def __enable_scan_buttons(self): - self.refresh_basic_stats_button.configure(state="normal") - self.refresh_advanced_stats_button.configure(state="normal") - - def __cancel_concept_stats(self): - self.cancel_scan_flag.set() - - def __auto_update_concept_stats(self): - try: - self.__update_concept_stats() #load stats from config if available, else raises KeyError - if self.concept.concept_stats["file_size"] == 0: #force rescan if empty - raise KeyError - except KeyError: - concept_path = self.get_concept_path(self.concept.path) - if concept_path: - self.__get_concept_stats(False, 2) #force rescan if config is empty, timeout of 2 sec - if self.concept.concept_stats["processing_time"] < 0.1: - self.__get_concept_stats(True, 2) #do advanced scan automatically if basic took <0.1s - def destroy(self): if self.bucket_fig is not None: plt.close(self.bucket_fig) @@ -930,5 +169,5 @@ def destroy(self): super().destroy() - def __ok(self): + def _ok(self): self.destroy() diff --git a/modules/ui/CtkConfigListView.py b/modules/ui/CtkConfigListView.py index 75d69252a..72995bfcc 100644 --- a/modules/ui/CtkConfigListView.py +++ b/modules/ui/CtkConfigListView.py @@ -1,27 +1,19 @@ import contextlib -import copy -import json -import os -import tkinter as tk -from abc import ABCMeta, abstractmethod +from abc import ABC -from modules.util import path_util -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.UIState import UIState +from modules.ui.BaseConfigListView import BaseConfigListView +from modules.util.ui import ctk_components, dialogs import customtkinter as ctk -class ConfigList(metaclass=ABCMeta): +class CtkConfigListView(BaseConfigListView, ABC): def __init__( self, master, - train_config: TrainConfig, - ui_state: UIState, + controller, + ui_state, from_external_file: bool, attr_name: str = "", enable_key: str = "enabled", @@ -29,326 +21,51 @@ def __init__( default_config_name: str = "", add_button_text: str = "", add_button_tooltip: str = "", - is_full_width: bool = "", + is_full_width: bool = False, show_toggle_button: bool = False, ): - self.master = master - self.train_config = train_config - self.ui_state = ui_state - self.from_external_file = from_external_file - self.attr_name = attr_name - self.enable_key = enable_key - - self.config_dir = config_dir - self.default_config_name = default_config_name - - self.is_full_width = is_full_width - - # From search-concepts - self.filters = {"search": "", "type": "ALL", "show_disabled": True} - self.widgets_initialized = False - - # From master - self.toggle_button = None - self.show_toggle_button = show_toggle_button - self.is_opening_window = False - self._is_current_item_enabled = False - - self.master.grid_rowconfigure(0, weight=0) - self.master.grid_rowconfigure(1, weight=1) - self.master.grid_columnconfigure(0, weight=1) - - if self.from_external_file: - self.top_frame = ctk.CTkFrame(self.master, fg_color="transparent") - self.top_frame.grid(row=0, column=0, sticky="nsew") - - self.configs_dropdown = None - self.element_list = None - - self.configs = [] - self.__load_available_config_names() - - 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) - 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.element_list = None - 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() - - - - @abstractmethod - def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - pass - - @abstractmethod - def create_new_element(self) -> BaseConfig: - pass - - @abstractmethod - def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: - pass - - def _refresh_show_disabled_text(self): - return - - def _reset_filters(self): # pragma: no cover - default noop - search_var = getattr(self, 'search_var', None) - filter_var = getattr(self, 'filter_var', None) - show_disabled_var = getattr(self, 'show_disabled_var', None) - - if search_var: - search_var.set("") - if filter_var: - filter_var.set("ALL") - if show_disabled_var: - show_disabled_var.set(True) - if search_var and hasattr(self, '_update_filters'): - self._update_filters() - - def _update_item_enabled_state(self): - # Only count items that match current filters - self._is_current_item_enabled = any( - item.ui_state.get_var(self.enable_key).get() - for i, item in enumerate(self.widgets) - if i < len(self.current_config) and self._element_matches_filters(self.current_config[i]) + BaseConfigListView.__init__(self, ctk_components) + + master.grid_rowconfigure(0, weight=0) + master.grid_rowconfigure(1, weight=1) + master.grid_columnconfigure(0, weight=1) + + self.build( + master, controller, ui_state, from_external_file, + attr_name=attr_name, + enable_key=enable_key, + config_dir=config_dir, + default_config_name=default_config_name, + add_button_text=add_button_text, + add_button_tooltip=add_button_tooltip, + is_full_width=is_full_width, + show_toggle_button=show_toggle_button, ) - def _update_toggle_button_text(self): - if not self.show_toggle_button: - return - self._update_item_enabled_state() - if self.toggle_button is not None: - self.toggle_button.configure(text="Disable" if self._is_current_item_enabled else "Enable") - - def _toggle(self): - self._toggle_items() - - def _toggle_items(self): - enable_state = not self._is_current_item_enabled - - # Only toggle items that match current filters - for i, widget in enumerate(self.widgets): - if i < len(self.current_config) and self._element_matches_filters(self.current_config[i]): - widget.ui_state.get_var(self.enable_key).set(enable_state) - self.save_current_config() - - self._update_widget_visibility() - - def __create_configs_dropdown(self): - if self.configs_dropdown is not None: - self.configs_dropdown.destroy() - - self.configs_dropdown = components.options_kv( - self.top_frame, 0, 0, self.configs, self.ui_state, self.attr_name, self.__load_current_config - ) - self._update_toggle_button_text() - - def _create_element_list(self, **filters): - if not self.from_external_file: - self.current_config = getattr(self.train_config, self.attr_name) - - self.filters.update(filters) - - if not self.widgets_initialized: - self._initialize_all_widgets() - self.widgets_initialized = True - - self._update_widget_visibility() - self._update_toggle_button_text() - - def _initialize_all_widgets(self): - self.widgets = [] - if self.element_list is not None: - self.element_list.destroy() - - self.element_list = ctk.CTkScrollableFrame(self.master, fg_color="transparent") - self.element_list.grid(row=1, column=0, sticky="nsew") + def _create_top_frame(self, master): + frame = ctk.CTkFrame(master, fg_color="transparent") + frame.grid(row=0, column=0, sticky="nsew") + return frame + def _create_element_list_frame(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid(row=1, column=0, sticky="nsew") if self.is_full_width: - self.element_list.grid_columnconfigure(0, weight=1) - - for i, element in enumerate(self.current_config): - widget = self.create_widget( - self.element_list, element, i, - self.__open_element_window, - self.__remove_element, - self.__clone_element, - self.save_current_config - ) - self.widgets.append(widget) - - def _update_widget_visibility(self): - visible_index = 0 - - for i, widget in enumerate(self.widgets): - if i < len(self.current_config): - element = self.current_config[i] - - if self._element_matches_filters(element): - widget.visible_index = visible_index - widget.place_in_list() - visible_index += 1 - else: - widget.grid_remove() - - def __load_available_config_names(self): - if os.path.isdir(self.config_dir): - for path in os.listdir(self.config_dir): - path = path_util.canonical_join(self.config_dir, path) - if path.endswith(".json") and os.path.isfile(path): - name = os.path.basename(path) - name = os.path.splitext(name)[0] - self.configs.append((name, path)) - - if len(self.configs) == 0: - name = self.default_config_name.removesuffix(".json") - self.__create_config(name) - self.save_current_config() - - def __create_config(self, name: str): - name = path_util.safe_filename(name) - path = path_util.canonical_join(self.config_dir, f"{name}.json") - self.configs.append((name, path)) - self.__create_configs_dropdown() - - def __add_config(self): - dialogs.StringInputDialog(self.master, "name", "Name", self.__create_config) - - def __add_element(self): - new_element = self.create_new_element() - 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 - widget = self.create_widget( - self.element_list, new_element, i, - self.__open_element_window, - self.__remove_element, - self.__clone_element, - self.save_current_config - ) - self.widgets.append(widget) - self._update_widget_visibility() - else: - self.widgets_initialized = False - self._create_element_list() - self.save_current_config() - - def __clone_element(self, clone_i, modify_element_fun=None): - new_element = copy.deepcopy(self.current_config[clone_i]) - - if modify_element_fun is not None: - new_element = modify_element_fun(new_element) - self.current_config.append(new_element) - if self.widgets_initialized and self.element_list is not None: - i = len(self.current_config) - 1 - widget = self.create_widget( - self.element_list, new_element, i, - self.__open_element_window, - self.__remove_element, - self.__clone_element, - self.save_current_config - ) - self.widgets.append(widget) - self._update_widget_visibility() - else: - self.widgets_initialized = False - self._create_element_list() - self.save_current_config() - - 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() - else: - self.widgets_initialized = False - self._create_element_list() - self.save_current_config() - - def __load_current_config(self, filename): - try: - with open(filename, "r") as f: - self.current_config = [] - - loaded_config_json = json.load(f) - for element_json in loaded_config_json: - element = self.create_new_element().from_dict(element_json) - self.current_config.append(element) - except (FileNotFoundError, json.JSONDecodeError) as e: - print(f"Failed to load config from {filename}: {e}") - self.current_config = [] - - # reset filters when switching configs - if hasattr(self, '_reset_filters') and self.widgets_initialized: - self._reset_filters() - - self.widgets_initialized = False - self._create_element_list() - self._update_toggle_button_text() - - def save_current_config(self): - if self.from_external_file: - try: - if not os.path.exists(self.config_dir): - os.makedirs(self.config_dir, exist_ok=True) - - write_json_atomic( - getattr(self.train_config, self.attr_name), - [element.to_dict() for element in self.current_config] - ) - except (OSError) as e: - print(f"Failed to save config: {e}") + frame.grid_columnconfigure(0, weight=1) + return frame - self._update_toggle_button_text() + def _wait_for_window(self, window): + self.master.wait_window(window) - if self.widgets_initialized: - try: - self._update_widget_visibility() - except (tk.TclError, AttributeError) as e: - print.debug(f"Widget visibility update failed: {e}") + def _remove_widget_from_layout(self, widget): + widget.grid_remove() - # let subclass refresh any show-disabled UI - if hasattr(self, '_refresh_show_disabled_text'): - self._refresh_show_disabled_text() + def _destroy_widget(self, widget): + with contextlib.suppress(AttributeError): + widget.destroy() - def _element_matches_filters(self, element): - return True # Show all by default + def _destroy_frame(self, frame): + frame.destroy() - def __open_element_window(self, i, ui_state): - if self.is_opening_window: - return - self.is_opening_window = True - try: - window = self.open_element_window(i, ui_state) - self.master.wait_window(window) - try: - if self.widgets is not None and 0 <= i < len(self.widgets): - self.widgets[i].configure_element() - except Exception: - self.widgets_initialized = False - self._create_element_list() - self.save_current_config() - finally: - self.is_opening_window = False + def _show_name_dialog(self, callback): + dialogs.StringInputDialog(self.master, "name", "Name", callback) diff --git a/modules/ui/CtkConvertModelUIView.py b/modules/ui/CtkConvertModelUIView.py index 6cb1b507a..782637348 100644 --- a/modules/ui/CtkConvertModelUIView.py +++ b/modules/ui/CtkConvertModelUIView.py @@ -1,33 +1,18 @@ -import traceback -from uuid import uuid4 - -from modules.util import create -from modules.util.args.ConvertModelArgs import ConvertModelArgs -from modules.util.config.TrainConfig import QuantizationConfig -from modules.util.enum.DataType import DataType -from modules.util.enum.ModelFormat import ModelFormat -from modules.util.enum.ModelType import ModelType -from modules.util.enum.PathIOType import PathIOType -from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.ModelNames import EmbeddingName, ModelNames -from modules.util.torch_util import torch_gc -from modules.util.ui import components +from modules.ui.BaseConvertModelUIView import BaseConvertModelUIView +from modules.ui.ConvertModelUIController import ConvertModelUIController +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState import customtkinter as ctk -class ConvertModelUI(ctk.CTkToplevel): - def __init__(self, parent, *args, **kwargs): - super().__init__(parent, *args, **kwargs) - self.parent = parent - - self.parent = parent - self.convert_model_args = ConvertModelArgs.default_values() - self.ui_state = UIState(self, self.convert_model_args) - self.button = None +class CtkConvertModelUIView(BaseConvertModelUIView, ctk.CTkToplevel): + def __init__(self, parent, controller: ConvertModelUIController, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseConvertModelUIView.__init__(self, ctk_components) + ui_state = CtkUIState(self, controller.convert_model_args) self.title("Convert models") self.geometry("550x350") @@ -38,133 +23,12 @@ def __init__(self, parent, *args, **kwargs): self.frame.grid_columnconfigure(0, weight=0) self.frame.grid_columnconfigure(1, weight=1) - self.main_frame(self.frame) + self.build_content(self.frame, controller, ui_state) self.frame.pack(fill="both", expand=True) self.wait_visibility() self.focus_set() self.after(200, lambda: set_window_icon(self)) - - def main_frame(self, master): - # model type - components.label(master, 0, 0, "Model Type", - tooltip="Type of the model") - components.options_kv(master, 0, 1, [ #TODO simplify - ("Stable Diffusion 1.5", ModelType.STABLE_DIFFUSION_15), - ("Stable Diffusion 1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), - ("Stable Diffusion 2.0", ModelType.STABLE_DIFFUSION_20), - ("Stable Diffusion 2.0 Inpainting", ModelType.STABLE_DIFFUSION_20_INPAINTING), - ("Stable Diffusion 2.1", ModelType.STABLE_DIFFUSION_21), - ("Stable Diffusion 3", ModelType.STABLE_DIFFUSION_3), - ("Stable Diffusion 3.5", ModelType.STABLE_DIFFUSION_35), - ("Stable Diffusion XL 1.0 Base", ModelType.STABLE_DIFFUSION_XL_10_BASE), - ("Stable Diffusion XL 1.0 Base Inpainting", ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING), - ("Wuerstchen v2", ModelType.WUERSTCHEN_2), - ("Stable Cascade", ModelType.STABLE_CASCADE_1), - ("PixArt Alpha", ModelType.PIXART_ALPHA), - ("PixArt Sigma", ModelType.PIXART_SIGMA), - ("Flux Dev", ModelType.FLUX_DEV_1), - ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), - ("Flux 2", ModelType.FLUX_2), - ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), - ("Chroma1", ModelType.CHROMA_1), #TODO does this just work? HiDream is not here - ("QwenImage", ModelType.QWEN), #TODO does this just work? HiDream is not here - ("ZImage", ModelType.Z_IMAGE), - ], self.ui_state, "model_type") - - # training method - components.label(master, 1, 0, "Model Type", - tooltip="The type of model to convert") - components.options_kv(master, 1, 1, [ - ("Base Model", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ("Embedding", TrainingMethod.EMBEDDING), - ], self.ui_state, "training_method") - - # input name - components.label(master, 2, 0, "Input name", - tooltip="Filename, directory or hugging face repository of the base model") - components.path_entry( - master, 2, 1, self.ui_state, "input_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # output data type - components.label(master, 3, 0, "Output Data Type", - tooltip="Precision to use when saving the output model") - components.options_kv(master, 3, 1, [ - ("float32", DataType.FLOAT_32), - ("float16", DataType.FLOAT_16), - ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "output_dtype") - - # output format - components.label(master, 4, 0, "Output Format", - tooltip="Format to use when saving the output model") - components.options_kv(master, 4, 1, [ - ("Safetensors", ModelFormat.SAFETENSORS), - ("Diffusers", ModelFormat.DIFFUSERS), - ], self.ui_state, "output_model_format") - - # output model destination - components.label(master, 5, 0, "Model Output Destination", - tooltip="Filename or directory where the output model is saved") - components.path_entry( - master, 5, 1, self.ui_state, "output_model_destination", - mode="file", - io_type=PathIOType.MODEL, - ) - - self.button = components.button(master, 6, 1, "Convert", self.convert_model) - - def convert_model(self): - try: - self.button.configure(state="disabled") - model_loader = create.create_model_loader( - model_type=self.convert_model_args.model_type, - training_method=self.convert_model_args.training_method - ) - model_saver = create.create_model_saver( - model_type=self.convert_model_args.model_type, - training_method=self.convert_model_args.training_method - ) - - print("Loading model " + self.convert_model_args.input_name) - if self.convert_model_args.training_method in [TrainingMethod.FINE_TUNE]: - model = model_loader.load( - model_type=self.convert_model_args.model_type, - model_names=ModelNames( - base_model=self.convert_model_args.input_name, - ), - weight_dtypes=self.convert_model_args.weight_dtypes(), - quantization=QuantizationConfig.default_values(), - ) - elif self.convert_model_args.training_method in [TrainingMethod.LORA, TrainingMethod.EMBEDDING]: - model = model_loader.load( - model_type=self.convert_model_args.model_type, - model_names=ModelNames( - base_model=None, - lora=self.convert_model_args.input_name, - embedding=EmbeddingName(str(uuid4()), self.convert_model_args.input_name), - ), - weight_dtypes=self.convert_model_args.weight_dtypes(), - quantization=QuantizationConfig.default_values(), - ) - else: - raise Exception("could not load model: " + self.convert_model_args.input_name) - - print("Saving model " + self.convert_model_args.output_model_destination) - model_saver.save( - model=model, - model_type=self.convert_model_args.model_type, - output_model_format=self.convert_model_args.output_model_format, - output_model_destination=self.convert_model_args.output_model_destination, - dtype=self.convert_model_args.output_dtype.torch_dtype(), - ) - print("Model converted") - except Exception: - traceback.print_exc() - - torch_gc() - self.button.configure(state="normal") + def set_converting(self, active): + self.button.configure(state="disabled" if active else "normal") diff --git a/modules/ui/CtkGenerateCaptionsWindowView.py b/modules/ui/CtkGenerateCaptionsWindowView.py index 1690879f1..09d82f74b 100644 --- a/modules/ui/CtkGenerateCaptionsWindowView.py +++ b/modules/ui/CtkGenerateCaptionsWindowView.py @@ -2,27 +2,22 @@ import tkinter as tk from tkinter import filedialog +from modules.ui.BaseGenerateCaptionsWindowView import BaseGenerateCaptionsWindowView +from modules.ui.GenerateCaptionsWindowController import GenerateCaptionsWindowController from modules.util.ui.ui_utils import set_window_icon import customtkinter as ctk -class GenerateCaptionsWindow(ctk.CTkToplevel): - def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs): - """ - Window for generating captions for a folder of images - - Parameters: - parent (`Tk`): the parent window - path (`str`): the path to the folder - parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox - """ - super().__init__(parent, *args, **kwargs) - self.parent = parent +class CtkGenerateCaptionsWindowView(BaseGenerateCaptionsWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: GenerateCaptionsWindowController, path, parent_include_subdirectories, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) if path is None: path = "" + self.controller = controller + self.mode_var = ctk.StringVar(self, "Create if absent") self.modes = ["Replace all captions", "Create if absent", "Add as new line"] self.model_var = ctk.StringVar(self, "Blip") @@ -79,7 +74,7 @@ def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs) self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) self.progress.grid(row=7, column=1, sticky="w", padx=5, pady=5) - self.create_captions_button = ctk.CTkButton(self.frame, text="Create Captions", width=310, command=self.create_captions) + self.create_captions_button = ctk.CTkButton(self.frame, text="Create Captions", width=310, command=self._on_create_captions) self.create_captions_button.grid(row=8, column=0, columnspan=2, sticky="w", padx=5, pady=5) self.frame.pack(fill="both", expand=True) @@ -89,7 +84,6 @@ def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs) self.focus_set() self.after(200, lambda: set_window_icon(self)) - def browse_for_path(self, entry_box): # get the path from the user path = filedialog.askdirectory() @@ -106,25 +100,16 @@ def set_progress(self, value, max_value): self.progress_label.configure(text=f"{value}/{max_value}") self.progress.update() - def create_captions(self): - self.parent.load_captioning_model(self.model_var.get()) - - mode = { - "Replace all captions": "replace", - "Create if absent": "fill", - "Add as new line": "add", - }[self.mode_var.get()] - - self.parent.captioning_model.caption_folder( - sample_dir=self.path_entry.get(), + def _on_create_captions(self): + self.controller.create_captions( + model_name=self.model_var.get(), + path=self.path_entry.get(), initial_caption=self.caption_entry.get(), caption_prefix=self.prefix_entry.get(), caption_postfix=self.postfix_entry.get(), - mode=mode, - progress_callback=self.set_progress, + mode_str=self.mode_var.get(), include_subdirectories=self.include_subdirectories_var.get(), ) - self.parent.load_image() def destroy(self): with contextlib.suppress(tk.TclError): diff --git a/modules/ui/CtkGenerateMasksWindowView.py b/modules/ui/CtkGenerateMasksWindowView.py index daff0d3d5..631179fac 100644 --- a/modules/ui/CtkGenerateMasksWindowView.py +++ b/modules/ui/CtkGenerateMasksWindowView.py @@ -2,13 +2,15 @@ import tkinter as tk from tkinter import filedialog +from modules.ui.BaseGenerateMasksWindowView import BaseGenerateMasksWindowView +from modules.ui.GenerateMasksWindowController import GenerateMasksWindowController from modules.util.ui.ui_utils import set_window_icon import customtkinter as ctk -class GenerateMasksWindow(ctk.CTkToplevel): - def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs): +class CtkGenerateMasksWindowView(BaseGenerateMasksWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: GenerateMasksWindowController, path, parent_include_subdirectories, *args, **kwargs): """ Window for generating masks for a folder of images @@ -17,9 +19,9 @@ def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs) path (`str`): the path to the folder parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox """ - super().__init__(parent, *args, **kwargs) + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) - self.parent = parent + self.controller = controller if path is None: path = "" @@ -93,7 +95,7 @@ def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs) self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) self.progress.grid(row=9, column=1, sticky="w", padx=5, pady=5) - self.create_masks_button = ctk.CTkButton(self.frame, text="Create Masks", width=310, command=self.create_masks) + self.create_masks_button = ctk.CTkButton(self.frame, text="Create Masks", width=310, command=self._on_create_masks) self.create_masks_button.grid(row=10, column=0, columnspan=2, sticky="w", padx=5, pady=5) self.frame.pack(fill="both", expand=True) @@ -103,7 +105,6 @@ def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs) self.focus_set() self.after(200, lambda: set_window_icon(self)) - def browse_for_path(self, entry_box): # get the path from the user path = filedialog.askdirectory() @@ -120,29 +121,18 @@ def set_progress(self, value, max_value): self.progress_label.configure(text=f"{value}/{max_value}") self.progress.update() - def create_masks(self): - self.parent.load_masking_model(self.model_var.get()) - - mode = { - "Replace all masks": "replace", - "Create if absent": "fill", - "Add to existing": "add", - "Subtract from existing": "subtract", - "Blend with existing": "blend", - }[self.mode_var.get()] - - self.parent.masking_model.mask_folder( - sample_dir=self.path_entry.get(), - prompts=[self.prompt_entry.get()], - mode=mode, - alpha=float(self.alpha_entry.get()), - threshold=float(self.threshold_entry.get()), - smooth_pixels=int(self.smooth_entry.get()), - expand_pixels=int(self.expand_entry.get()), - progress_callback=self.set_progress, + def _on_create_masks(self): + self.controller.create_masks( + model_name=self.model_var.get(), + path=self.path_entry.get(), + prompt=self.prompt_entry.get(), + mode_str=self.mode_var.get(), + alpha_str=self.alpha_entry.get(), + threshold_str=self.threshold_entry.get(), + smooth_str=self.smooth_entry.get(), + expand_str=self.expand_entry.get(), include_subdirectories=self.include_subdirectories_var.get(), ) - self.parent.load_image() def destroy(self): with contextlib.suppress(tk.TclError): diff --git a/modules/ui/CtkLoraTabView.py b/modules/ui/CtkLoraTabView.py index 1c73d90ce..8caa1f171 100644 --- a/modules/ui/CtkLoraTabView.py +++ b/modules/ui/CtkLoraTabView.py @@ -1,25 +1,20 @@ -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.DataType import DataType +from modules.ui.BaseLoraTabView import BaseLoraTabView +from modules.ui.LoraTabController import LoraTabController from modules.util.enum.ModelType import PeftType -from modules.util.ui import components -from modules.util.ui.UIState import UIState -from modules.util.ui.validation_helpers import check_range +from modules.util.ui import ctk_components import customtkinter as ctk -class LoraTab: - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__() - +class CtkLoraTabView(BaseLoraTabView): + def __init__(self, master, controller: LoraTabController, ui_state): + BaseLoraTabView.__init__(self, ctk_components) self.master = master - self.train_config = train_config + self.controller = controller self.ui_state = ui_state - self.scroll_frame = None self.options_frame = None - self.refresh_ui() def refresh_ui(self): @@ -27,128 +22,20 @@ def refresh_ui(self): self.scroll_frame.destroy() self.scroll_frame = ctk.CTkFrame(self.master, fg_color="transparent") self.scroll_frame.grid(row=0, column=0, sticky="nsew") - self.scroll_frame.grid_columnconfigure(0, weight=0) self.scroll_frame.grid_columnconfigure(1, weight=1) self.scroll_frame.grid_columnconfigure(2, weight=2) - - components.label(self.scroll_frame, 0, 0, "Type", - tooltip="The type of low-parameter finetuning method.") - # This will instantly call self.setup_lora. - components.options_kv(self.scroll_frame, 0, 1, [ - ("LoRA", PeftType.LORA), - ("LoHa", PeftType.LOHA), - ("OFT v2", PeftType.OFT_2), - ], self.ui_state, "peft_type", command=self.setup_lora) + self.build(self.scroll_frame, self.controller, self.ui_state, self.setup_lora) def setup_lora(self, peft_type: PeftType): - if peft_type == PeftType.LOHA: - name = "LoHa" - elif peft_type == PeftType.OFT_2: - name = "OFT v2" - else: - name = "LoRA" - if self.options_frame: self.options_frame.destroy() self.options_frame = ctk.CTkFrame(self.scroll_frame, fg_color="transparent") self.options_frame.grid(row=1, column=0, columnspan=3, sticky="nsew") master = self.options_frame - master.grid_columnconfigure(0, weight=0, uniform="a") master.grid_columnconfigure(1, weight=1, uniform="a") master.grid_columnconfigure(2, minsize=50, uniform="a") master.grid_columnconfigure(3, weight=0, uniform="a") master.grid_columnconfigure(4, weight=1, uniform="a") - - # 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.path_entry( - master, 0, 1, self.ui_state, "lora_model_name", - mode="file", path_modifier=components.json_path_modifier - ) - entry.grid(row=0, column=1, columnspan=4) - - - # LoRA decomposition - if peft_type == PeftType.LORA: - components.label(master, 1, 3, "Decompose Weights (DoRA)", - tooltip="Decompose LoRA Weights (aka, DoRA).") - components.switch(master, 1, 4, self.ui_state, "lora_decompose") - - components.label(master, 2, 3, "Use Norm Epsilon (DoRA Only)", - tooltip="Add an epsilon to the norm divison calculation in DoRA. Can aid in training stability, and also acts as regularization.") - components.switch(master, 2, 4, self.ui_state, "lora_decompose_norm_epsilon") - components.label(master, 3, 3, "Apply on output axis (DoRA Only)", - tooltip="Apply the weight decomposition on the output axis instead of the input axis.") - components.switch(master, 3, 4, self.ui_state, "lora_decompose_output_axis") - - # LoRA and LoHA shared settings - if peft_type == PeftType.LORA or peft_type == PeftType.LOHA: - # rank - components.label(master, 1, 0, f"{name} rank", - tooltip=f"The rank parameter used when creating a new {name}") - components.entry(master, 1, 1, self.ui_state, "lora_rank", required=True, extra_validate=check_range(lower=1, message="Rank must be at least 1")) - - # alpha - components.label(master, 2, 0, f"{name} alpha", - tooltip=f"The alpha parameter used when creating a new {name}") - components.entry(master, 2, 1, self.ui_state, "lora_alpha", required=True) - - # Dropout Percentage - components.label(master, 3, 0, "Dropout Probability", - tooltip="Dropout probability. This percentage of model nodes will be randomly ignored at each training step. Helps with overfitting. 0 disables, 1 maximum.") - components.entry(master, 3, 1, self.ui_state, "dropout_probability") - - # weight dtype - components.label(master, 4, 0, f"{name} Weight Data Type", - tooltip=f"The {name} weight data type used for training. This can reduce memory consumption, but reduces precision") - components.options_kv(master, 4, 1, [ - ("float32", DataType.FLOAT_32), - ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "lora_weight_dtype") - - # For use with additional embeddings. - components.label(master, 5, 0, "Bundle Embeddings", - tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") - components.switch(master, 5, 1, self.ui_state, "bundle_additional_embeddings") - - # OFTv2 - elif peft_type == PeftType.OFT_2: - # Block Size - components.label(master, 1, 0, f"{name} Block Size", - tooltip=f"The block size parameter used when creating a new {name}") - components.entry(master, 1, 1, self.ui_state, "oft_block_size", required=True) - - # COFT - components.label(master, 1, 3, "Constrained OFT (COFT)", - tooltip="Use the constrained variant of OFT. This constrains the learned rotation to stay very close to the identity matrix, limiting adaptation to only small changes. This improves training stability, helps prevent overfitting on small datasets, and better preserves the base models original knowledge but it may lack expressiveness for tasks requiring substantial adaptation and introduces an additional hyperparameter (COFT Epsilon) that needs tuning.") - components.switch(master, 1, 4, self.ui_state, "oft_coft") - - components.label(master, 2, 3, "COFT Epsilon", - tooltip="The control strength of COFT. Only has an effect if COFT is enabled.") - components.entry(master, 2, 4, self.ui_state, "coft_eps") - - # Block Share - components.label(master, 3, 3, "Block Share", - tooltip="Share the OFT parameters between blocks. A single rotation matrix is shared across all blocks within a layer, drastically cutting the number of trainable parameters and yielding very compact adapter files, potentially improving generalization but at the cost of significant expressiveness, which can lead to underfitting on more complex or diverse tasks.") - components.switch(master, 3, 4, self.ui_state, "oft_block_share") - - # Dropout Percentage - components.label(master, 2, 0, "Dropout Probability", - tooltip="Dropout probability. This percentage of the rotated adapter nodes that will be randomly restored to the base model initial statue. Helps with overfitting. 0 disables, 1 maximum.") - components.entry(master, 2, 1, self.ui_state, "dropout_probability") - - # OFT weight dtype - components.label(master, 3, 0, f"{name} Weight Data Type", - tooltip=f"The {name} weight data type used for training. This can reduce memory consumption, but reduces precision") - components.options_kv(master, 3, 1, [ - ("float32", DataType.FLOAT_32), - ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "lora_weight_dtype") - - # For use with additional embeddings. - components.label(master, 4, 0, "Bundle Embeddings", - tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") - components.switch(master, 4, 1, self.ui_state, "bundle_additional_embeddings") + self.build_lora_options(master, self.controller, self.ui_state, peft_type) diff --git a/modules/ui/CtkModelTabView.py b/modules/ui/CtkModelTabView.py index ff17ea3ba..5b43b7dca 100644 --- a/modules/ui/CtkModelTabView.py +++ b/modules/ui/CtkModelTabView.py @@ -1,24 +1,17 @@ -from modules.util import create -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.ConfigPart import ConfigPart -from modules.util.enum.DataType import DataType -from modules.util.enum.ModelFormat import ModelFormat -from modules.util.enum.PathIOType import PathIOType -from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.ui import components -from modules.util.ui.UIState import UIState - -import customtkinter as ctk +from modules.ui.BaseModelTabView import BaseModelTabView +from modules.ui.ModelTabController import ModelTabController +from modules.util.ui import ctk_components -class ModelTab: +import customtkinter as ctk - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__() +class CtkModelTabView(BaseModelTabView): + def __init__(self, master, controller: ModelTabController, ui_state): + BaseModelTabView.__init__(self, ctk_components) self.master = master - self.train_config = train_config + self.controller = controller self.ui_state = ui_state master.grid_rowconfigure(0, weight=1) @@ -28,6 +21,13 @@ def __init__(self, master, train_config: TrainConfig, ui_state: UIState): self.refresh_ui() + def _make_svd_frames(self, parent, row: int): + svd_label_frame = ctk.CTkFrame(parent, fg_color="transparent") + svd_label_frame.grid(row=row, column=3, sticky="nsew") + svd_entry_frame = ctk.CTkFrame(parent, fg_color="transparent") + svd_entry_frame.grid(row=row, column=4, sticky="nsew") + return svd_label_frame, svd_entry_frame + def refresh_ui(self): if self.scroll_frame: self.scroll_frame.destroy() @@ -40,649 +40,9 @@ def refresh_ui(self): base_frame.grid(row=0, column=0, padx=5, pady=5, sticky="nsew") base_frame.grid_columnconfigure(0, weight=0) - base_frame.grid_columnconfigure(1, weight=10)#, minsize=500) + base_frame.grid_columnconfigure(1, weight=10) # , minsize=500) base_frame.grid_columnconfigure(2, minsize=50) base_frame.grid_columnconfigure(3, weight=0) base_frame.grid_columnconfigure(4, weight=1) - if self.train_config.model_type.is_stable_diffusion(): #TODO simplify - self.__setup_stable_diffusion_ui(base_frame) - if self.train_config.model_type.is_stable_diffusion_3(): - self.__setup_stable_diffusion_3_ui(base_frame) - elif self.train_config.model_type.is_stable_diffusion_xl(): - self.__setup_stable_diffusion_xl_ui(base_frame) - elif self.train_config.model_type.is_wuerstchen(): - self.__setup_wuerstchen_ui(base_frame) - elif self.train_config.model_type.is_pixart(): - self.__setup_pixart_alpha_ui(base_frame) - elif self.train_config.model_type.is_flux_1(): - self.__setup_flux_ui(base_frame) - elif self.train_config.model_type.is_flux_2(): - self.__setup_flux_2_ui(base_frame) - elif self.train_config.model_type.is_z_image(): - self.__setup_z_image_ui(base_frame) - elif self.train_config.model_type.is_chroma(): - self.__setup_chroma_ui(base_frame) - elif self.train_config.model_type.is_qwen(): - self.__setup_qwen_ui(base_frame) - elif self.train_config.model_type.is_sana(): - self.__setup_sana_ui(base_frame) - elif self.train_config.model_type.is_hunyuan_video(): - self.__setup_hunyuan_video_ui(base_frame) - elif self.train_config.model_type.is_hi_dream(): - self.__setup_hi_dream_ui(base_frame) - elif self.train_config.model_type.is_ernie(): - self.__setup_ernie_ui(base_frame) - - def __setup_stable_diffusion_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_unet=True, - has_text_encoder=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method in [ - TrainingMethod.FINE_TUNE, - TrainingMethod.FINE_TUNE_VAE, - ], - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_stable_diffusion_3_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_text_encoder_3=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_flux_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_flux_2_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_z_image_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_ernie_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_chroma_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_qwen_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_stable_diffusion_xl_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_unet=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_wuerstchen_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_prior=True, - allow_override_prior=self.train_config.model_type.is_stable_cascade(), - has_text_encoder=True, - ) - row = self.__create_effnet_encoder_components(frame, row) - row = self.__create_decoder_components(frame, row, self.train_config.model_type.is_wuerstchen_v2()) - row = self.__create_output_components( - frame, - row, - allow_safetensors=self.train_config.training_method != TrainingMethod.FINE_TUNE - or self.train_config.model_type.is_stable_cascade(), - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_pixart_alpha_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - has_text_encoder=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_sana_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - has_text_encoder=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=self.train_config.training_method != TrainingMethod.FINE_TUNE, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_hunyuan_video_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_hi_dream_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_text_encoder_3=True, - has_text_encoder_4=True, - allow_override_text_encoder_4=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __create_dtype_options(self, include_gguf: bool=False, include_a8: bool=False) -> list[tuple[str, DataType]]: - options = [ - ("float32", DataType.FLOAT_32), - ("bfloat16", DataType.BFLOAT_16), - ("float16", DataType.FLOAT_16), - ("float8 (W8)", DataType.FLOAT_8), - # ("int8", DataType.INT_8), # TODO: reactivate when the int8 implementation is fixed in bitsandbytes: https://github.com/bitsandbytes-foundation/bitsandbytes/issues/1332 - ("nfloat4", DataType.NFLOAT_4), - ] - if include_a8: - options += [ - ("float W8A8", DataType.FLOAT_W8A8), - ("int W8A8", DataType.INT_W8A8), - ] - - if include_gguf: - options.append(("GGUF", DataType.GGUF)) - if include_a8: - options += [ - ("GGUF A8 float", DataType.GGUF_A8_FLOAT), - ("GGUF A8 int", DataType.GGUF_A8_INT), - ] - - return options - - def __create_base_dtype_components(self, frame, row: int) -> int: - # huggingface token - components.label(frame, row, 0, "Hugging Face Token", - tooltip="Enter your Hugging Face access token if you have used a protected Hugging Face repository below.\nThis value is stored separately, not saved to your configuration file. " - "Go to https://huggingface.co/settings/tokens to create an access token.", - wide_tooltip=True) - components.entry(frame, row, 1, self.ui_state, "secrets.huggingface_token") - - row += 1 - - # base model - components.label(frame, row, 0, "Base Model", - tooltip="Filename, directory or Hugging Face repository of the base model") - components.path_entry( - frame, row, 1, self.ui_state, "base_model_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # compile - components.label(frame, row, 3, "Compile transformer blocks", - tooltip="Uses torch.compile and Triton to significantly speed up training. Only applies to transformer/unet. Disable in case of compatibility issues.") - components.switch(frame, row, 4, self.ui_state, "compile") - - row += 1 - - return row - - def __create_base_components( - self, - frame, - row: int, - has_unet: bool = False, - has_prior: bool = False, - allow_override_prior: bool = False, - has_transformer: bool = False, - allow_override_transformer: bool = False, - allow_override_text_encoder_4: bool = False, - has_text_encoder: bool = False, - has_text_encoder_1: bool = False, - has_text_encoder_2: bool = False, - has_text_encoder_3: bool = False, - has_text_encoder_4: bool = False, - has_vae: bool = False, - ) -> int: - if has_unet: - # unet weight dtype - components.label(frame, row, 3, "UNet Data Type", - tooltip="The unet weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(include_a8=True), - self.ui_state, "unet.weight_dtype") - - row += 1 - - if has_prior: - if allow_override_prior: - # prior model - components.label(frame, row, 0, "Prior Model", - tooltip="Filename, directory or Hugging Face repository of the prior model") - components.path_entry( - frame, row, 1, self.ui_state, "prior.model_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # prior weight dtype - components.label(frame, row, 3, "Prior Data Type", - tooltip="The prior weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "prior.weight_dtype") - - row += 1 - - if has_transformer: - if allow_override_transformer: - # 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.path_entry( - frame, row, 1, self.ui_state, "transformer.model_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # transformer weight dtype - components.label(frame, row, 3, "Transformer Data Type", - tooltip="The transformer weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(include_gguf=True, include_a8=True), - self.ui_state, "transformer.weight_dtype") - - row += 1 - - cls = create.get_model_setup_class(self.train_config.model_type, self.train_config.training_method) - presets = cls.LAYER_PRESETS if cls is not None else {"full": []} - - components.label(frame, row, 0, "Quantization") - components.layer_filter_entry(frame, row, 1, self.ui_state, - preset_var_name="quantization.layer_filter_preset", presets=presets, - preset_label="Quantization Layer Filter", - preset_tooltip="Select a preset defining which layers to quantize. Quantization of certain layers can decrease model quality. Only applies to the transformer/unet", - entry_var_name="quantization.layer_filter", - entry_tooltip="Comma-separated list of layers to quantize. Regular expressions (if toggled) are supported. Any model layer with a matching name will be quantized", - regex_var_name="quantization.layer_filter_regex", - regex_tooltip="If enabled, layer filter patterns are interpreted as regular expressions. Otherwise, simple substring matching is used.", - frame_color="transparent", - ) - - # SVDQuant - create vertical grids to match the size of layer_filter_entry - svd_label_frame = ctk.CTkFrame(frame, fg_color="transparent") - svd_label_frame.grid(row=row, column=3, sticky="nsew") - svd_entry_frame = ctk.CTkFrame(frame, fg_color="transparent") - svd_entry_frame.grid(row=row, column=4, sticky="nsew") - components.label(svd_label_frame, 0, 0, "SVDQuant", - tooltip="What datatype to use for SVDQuant weights decomposition.") - components.options_kv(svd_entry_frame, 0, 0, [("disabled", DataType.NONE), ("float32", DataType.FLOAT_32), ("bfloat16", DataType.BFLOAT_16)], - self.ui_state, "quantization.svd_dtype") - components.label(svd_label_frame, 1, 0, "SVDQuant Rank", - tooltip="Rank for SVDQuant weights decomposition") - components.entry(svd_entry_frame, 1, 0, self.ui_state, "quantization.svd_rank") - row += 1 - - - if has_text_encoder: - # text encoder weight dtype - components.label(frame, row, 3, "Text Encoder Data Type", - tooltip="The text encoder weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "text_encoder.weight_dtype") - - row += 1 - - if has_text_encoder_1: - # text encoder 1 weight dtype - components.label(frame, row, 3, "Text Encoder 1 Data Type", - tooltip="The text encoder 1 weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "text_encoder.weight_dtype") - - row += 1 - - if has_text_encoder_2: - # text encoder 2 weight dtype - components.label(frame, row, 3, "Text Encoder 2 Data Type", - tooltip="The text encoder 2 weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "text_encoder_2.weight_dtype") - - row += 1 - - if has_text_encoder_3: - # text encoder 3 weight dtype - components.label(frame, row, 3, "Text Encoder 3 Data Type", - tooltip="The text encoder 3 weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "text_encoder_3.weight_dtype") - - row += 1 - - if has_text_encoder_4: - if allow_override_text_encoder_4: - # 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.path_entry( - frame, row, 1, self.ui_state, "text_encoder_4.model_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # text encoder 4 weight dtype - components.label(frame, row, 3, "Text Encoder 4 Data Type", - tooltip="The text encoder 4 weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "text_encoder_4.weight_dtype") - - row += 1 - - 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.path_entry( - frame, row, 1, self.ui_state, "vae.model_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # vae weight dtype - components.label(frame, row, 3, "VAE Data Type", - tooltip="The vae weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "vae.weight_dtype") - - row += 1 - - return row - - 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.path_entry( - frame, row, 1, self.ui_state, "effnet_encoder.model_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # effnet encoder weight dtype - components.label(frame, row, 3, "Effnet Encoder Data Type", - tooltip="The effnet encoder weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "effnet_encoder.weight_dtype") - - row += 1 - - return row - - def __create_decoder_components( - self, - frame, - row: int, - has_text_encoder: bool, - ) -> int: - # decoder model - components.label(frame, row, 0, "Decoder Model", - tooltip="Filename, directory or Hugging Face repository of the decoder model") - components.path_entry( - frame, row, 1, self.ui_state, "decoder.model_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # decoder weight dtype - components.label(frame, row, 3, "Decoder Data Type", - tooltip="The decoder weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "decoder.weight_dtype") - - row += 1 - - if has_text_encoder: - # decoder text encoder weight dtype - components.label(frame, row, 3, "Decoder Text Encoder Data Type", - tooltip="The decoder text encoder weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "decoder_text_encoder.weight_dtype") - - row += 1 - - # decoder vqgan weight dtype - components.label(frame, row, 3, "Decoder VQGAN Data Type", - tooltip="The decoder vqgan weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "decoder_vqgan.weight_dtype") - - row += 1 - - return row - - def __create_output_components( - self, - frame, - row: int, - allow_safetensors: bool = False, - allow_diffusers: bool = False, - allow_legacy_safetensors: bool = False, - allow_comfy: bool = False, - ) -> int: - # output model destination - components.label(frame, row, 0, "Model Output Destination", - tooltip="Filename or directory where the output model is saved") - components.path_entry( - frame, row, 1, self.ui_state, "output_model_destination", - mode="file", - io_type=PathIOType.MODEL, - ) - - # output data type - components.label(frame, row, 3, "Output Data Type", - tooltip="Precision to use when saving the output model") - components.options_kv(frame, row, 4, [ - ("float16", DataType.FLOAT_16), - ("float32", DataType.FLOAT_32), - ("bfloat16", DataType.BFLOAT_16), - ("float8", DataType.FLOAT_8), - ("nfloat4", DataType.NFLOAT_4), - ], self.ui_state, "output_dtype") - - row += 1 - - # output format - formats = [] - if allow_safetensors: - formats.append(("Safetensors", ModelFormat.SAFETENSORS)) - if allow_diffusers: - formats.append(("Diffusers", ModelFormat.DIFFUSERS)) - # if allow_legacy_safetensors: - # formats.append(("Legacy Safetensors", ModelFormat.LEGACY_SAFETENSORS)) - if allow_comfy: - formats.append(("Comfy LoRA", ModelFormat.COMFY_LORA)) - - 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") - - # include config - components.label(frame, row, 3, "Include Config", - tooltip="Include the training configuration in the final model. Only supported for safetensors files. " - "None: No config is included. " - "Settings: All training settings are included. " - "All: All settings, including the samples and concepts are included.") - components.options_kv(frame, row, 4, [ - ("None", ConfigPart.NONE), - ("Settings", ConfigPart.SETTINGS), - ("All", ConfigPart.ALL), - ], self.ui_state, "include_train_config") - - row += 1 - - return row + self.build_content(base_frame, self.controller, self.ui_state) diff --git a/modules/ui/CtkMuonAdamWindowView.py b/modules/ui/CtkMuonAdamWindowView.py index 5879ab432..3dc48ab74 100644 --- a/modules/ui/CtkMuonAdamWindowView.py +++ b/modules/ui/CtkMuonAdamWindowView.py @@ -1,42 +1,17 @@ -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.Optimizer import Optimizer -from modules.util.optimizer_util import OPTIMIZER_DEFAULT_PARAMETERS -from modules.util.ui import components +from modules.ui.BaseMuonAdamWindowView import BaseMuonAdamWindowView +from modules.ui.MuonAdamWindowController import MuonAdamWindowController +from modules.util.ui import ctk_components from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState import customtkinter as ctk -MUON_AUX_ADAM_DEFAULTS = { - "beta1": 0.9, - "beta2": 0.999, - "eps": 1e-8, - "weight_decay": 0.0, -} -class MuonAdamWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - train_config: TrainConfig, - ui_state: UIState, - parent_optimizer_type: Optimizer, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) - - self.parent = parent - self.train_config = train_config - self.adam_ui_state = ui_state - self.parent_optimizer_type = parent_optimizer_type - - if self.parent_optimizer_type == Optimizer.MUON: - self.title("Muon's Auxiliary AdamW Settings") - self.adam_params_def = MUON_AUX_ADAM_DEFAULTS - else: - self.title("Muon_adv's Auxiliary AdamW_adv Settings") - self.adam_params_def = OPTIMIZER_DEFAULT_PARAMETERS[Optimizer.ADAMW_ADV] +class CtkMuonAdamWindowView(BaseMuonAdamWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: MuonAdamWindowController, ui_state, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseMuonAdamWindowView.__init__(self, ctk_components) + self.title(controller.get_title()) self.geometry("800x500") self.resizable(True, True) @@ -44,64 +19,19 @@ def __init__( self.grid_rowconfigure(1, weight=0) self.grid_columnconfigure(0, weight=1) - self.frame = ctk.CTkScrollableFrame(self, fg_color="transparent") - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + frame = ctk.CTkScrollableFrame(self, fg_color="transparent") + frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) - self.frame.grid_columnconfigure(2, minsize=50) - self.frame.grid_columnconfigure(3, weight=0) - self.frame.grid_columnconfigure(4, weight=1) + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) - components.button(self, 1, 0, "ok", command=self.destroy) - self.create_adam_params_ui(self.frame) + self.components.button(self, 1, 0, "ok", command=self.destroy) + self.build_content(frame, controller, ui_state) self.wait_visibility() self.grab_set() self.focus_set() self.after(200, lambda: set_window_icon(self)) - - def create_adam_params_ui(self, master): - # This is a large map, copied from OptimizerParamsWindow for simplicity. - # @formatter:off - KEY_DETAIL_MAP = { - 'alpha': {'title': 'Alpha', 'tooltip': 'Smoothing parameter for RMSprop and others.', 'type': 'float'}, - 'beta1': {'title': 'Beta1', 'tooltip': 'optimizer_momentum term.', 'type': 'float'}, - 'beta2': {'title': 'Beta2', 'tooltip': 'Coefficients for computing running averages of gradient.', 'type': 'float'}, - 'eps': {'title': 'EPS', 'tooltip': 'A small value to prevent division by zero.', 'type': 'float'}, - 'stochastic_rounding': {'title': 'Stochastic Rounding', 'tooltip': 'Stochastic rounding for weight updates. Improves quality when using bfloat16 weights.', 'type': 'bool'}, - 'use_bias_correction': {'title': 'Bias Correction', 'tooltip': 'Turn on Adam\'s bias correction.', 'type': 'bool'}, - 'weight_decay': {'title': 'Weight Decay', 'tooltip': 'Regularization to prevent overfitting.', 'type': 'float'}, - 'use_orthograd': {'title': 'use_orthograd', 'tooltip': 'Use orthograd method', 'type': 'bool'}, - 'nnmf_factor': {'title': 'Factored Optimizer', 'tooltip': 'Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. This is an approximation which may slightly alter training dynamics.', 'type': 'bool'}, - 'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization.', 'type': 'bool'}, - 'use_atan2': {'title': 'Atan2 Scaling', 'tooltip': 'A robust replacement for eps, which also incorporates gradient clipping, bounding and stabilizing the optimizer updates.', 'type': 'bool'}, - 'use_AdEMAMix': {'title': 'AdEMAMix EMA', 'tooltip': 'Adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, and accelerate the training.', 'type': 'bool'}, - 'beta3_ema': {'title': 'Beta3 EMA', 'tooltip': 'Coefficient for slow-moving EMA of AdEMAMix.', 'type': 'float'}, - 'Simplified_AdEMAMix': {'title': 'Simplified AdEMAMix', 'tooltip': "Enables a simplified, single-EMA variant of AdEMAMix. Instead of blending two moving averages (fast and slow momentum), this version combines the raw current gradient (controlled by 'Grad α') directly with a single theory-based momentum. This makes the optimizer highly responsive to recent gradient information, which can accelerate training in all batch size scenarios when tuned correctly.", 'type': 'bool'}, - 'alpha_grad': {'title': 'Grad α', 'tooltip': 'Controls the mixing coefficient between raw gradients and momentum gradients in Simplified AdEMAMix. Higher values (e.g., 10-100) emphasize recent gradients, suitable for small batch sizes to reduce noise. Lower values (e.g., 0-1) emphasize historical gradients, suitable for large batch sizes for stability. Setting to 0 uses only momentum gradients without raw gradient contribution.', 'type': 'float'}, - 'kourkoutas_beta': {'title': 'Kourkoutas Beta', 'tooltip': 'Enables a layer-wise dynamic β₂ adaptation. This feature makes the optimizer more responsive to "spiky" gradients by lowering β₂ during periods of high variance, and more stable during calm periods by raising β₂ towards its maximum. It can significantly improve training stability and final loss.', 'type': 'bool'}, - } - # @formatter:on - - adam_params = self.adam_params_def - - for index, key in enumerate(adam_params.keys()): - if key not in KEY_DETAIL_MAP: - continue - - arg_info = KEY_DETAIL_MAP[key] - - title = arg_info['title'] - tooltip = arg_info['tooltip'] - param_type = arg_info['type'] - - row = index // 2 - col = 3 * (index % 2) - - components.label(master, row, col, title, tooltip=tooltip) - - if param_type != 'bool': - components.entry(master, row, col + 1, self.adam_ui_state, key) - else: - components.switch(master, row, col + 1, self.adam_ui_state, key) diff --git a/modules/ui/CtkOffloadingWindowView.py b/modules/ui/CtkOffloadingWindowView.py index 54035e121..b752f1602 100644 --- a/modules/ui/CtkOffloadingWindowView.py +++ b/modules/ui/CtkOffloadingWindowView.py @@ -1,75 +1,31 @@ -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.GradientCheckpointingMethod import ( - GradientCheckpointingMethod, -) -from modules.util.ui import components +from modules.ui.BaseOffloadingWindowView import BaseOffloadingWindowView +from modules.ui.OffloadingWindowController import OffloadingWindowController +from modules.util.ui import ctk_components from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState import customtkinter as ctk -class OffloadingWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - config: TrainConfig, - ui_state: UIState, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) - - self.config = config - self.ui_state = ui_state - self.image_preview_file_index = 0 - self.ax = None - self.canvas = None +class CtkOffloadingWindowView(BaseOffloadingWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: OffloadingWindowController, ui_state, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseOffloadingWindowView.__init__(self, ctk_components) self.title("Offloading") self.geometry("800x400") self.resizable(True, True) - self.grid_rowconfigure(0, weight=1) self.grid_columnconfigure(0, weight=1) - frame = self.__content_frame(self) + frame = ctk.CTkScrollableFrame(self, fg_color="transparent") + frame.grid_columnconfigure(0, weight=1) + frame.grid_columnconfigure(1, weight=1) + self.build_content(frame, controller, ui_state) + frame.pack(fill="both", expand=1) frame.grid(row=0, column=0, sticky='nsew') - components.button(self, 1, 0, "ok", self.__ok) + self.components.button(self, 1, 0, "ok", self.destroy) self.wait_visibility() self.grab_set() self.focus_set() self.after(200, lambda: set_window_icon(self)) - - - def __content_frame(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=1) - frame.grid_columnconfigure(1, weight=1) - - # timestep distribution - components.label(frame, 0, 0, "Gradient checkpointing", - tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") - components.options(frame, 0, 1, [str(x) for x in list(GradientCheckpointingMethod)], self.ui_state, - "gradient_checkpointing") - - # gradient checkpointing layer offloading - components.label(frame, 1, 0, "Async Offloading", - tooltip="Enables Asynchronous offloading.") - components.switch(frame, 1, 1, self.ui_state, "enable_async_offloading") - - # gradient checkpointing layer offloading - components.label(frame, 2, 0, "Offload Activations", - tooltip="Enables Activation Offloading") - components.switch(frame, 2, 1, self.ui_state, "enable_activation_offloading") - - # gradient checkpointing layer offloading - components.label(frame, 3, 0, "Layer offload fraction", - tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") - components.entry(frame, 3, 1, self.ui_state, "layer_offload_fraction") - - frame.pack(fill="both", expand=1) - return frame - - def __ok(self): - self.destroy() diff --git a/modules/ui/CtkOptimizerParamsWindowView.py b/modules/ui/CtkOptimizerParamsWindowView.py index 16063c26c..ecc1f6a38 100644 --- a/modules/ui/CtkOptimizerParamsWindowView.py +++ b/modules/ui/CtkOptimizerParamsWindowView.py @@ -1,38 +1,27 @@ import contextlib from tkinter import TclError -from modules.ui.MuonAdamWindow import MUON_AUX_ADAM_DEFAULTS, MuonAdamWindow -from modules.util.config.TrainConfig import TrainConfig, TrainOptimizerConfig -from modules.util.enum.Optimizer import Optimizer -from modules.util.optimizer_util import ( - OPTIMIZER_DEFAULT_PARAMETERS, - change_optimizer, - load_optimizer_defaults, - update_optimizer_config, -) -from modules.util.ui import components +from modules.ui.BaseOptimizerParamsWindowView import BaseOptimizerParamsWindowView +from modules.ui.CtkMuonAdamWindowView import CtkMuonAdamWindowView +from modules.ui.MuonAdamWindowController import MuonAdamWindowController +from modules.ui.OptimizerParamsWindowController import OptimizerParamsWindowController +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState import customtkinter as ctk -class OptimizerParamsWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - train_config: TrainConfig, - ui_state, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) +class CtkOptimizerParamsWindowView(BaseOptimizerParamsWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: OptimizerParamsWindowController, ui_state, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseOptimizerParamsWindowView.__init__(self, ctk_components) - self.parent = parent - self.train_config = train_config + self.controller = controller self.ui_state = ui_state self.optimizer_ui_state = ui_state.get_var("optimizer") - self.protocol("WM_DELETE_WINDOW", self.on_window_close) self.muon_adam_button = None + self.protocol("WM_DELETE_WINDOW", self.on_window_close) self.title("Optimizer Settings") self.geometry("800x500") @@ -51,206 +40,40 @@ def __init__( self.frame.grid_columnconfigure(3, weight=0) self.frame.grid_columnconfigure(4, weight=1) - components.button(self, 1, 0, "ok", command=self.on_window_close) - self.main_frame(self.frame) + self.components.button(self, 1, 0, "ok", command=self.on_window_close) + self.build_content(self.frame, controller, ui_state, self.optimizer_ui_state, + self.on_optimizer_change, self._load_defaults) + self._rebuild_dynamic_ui() self.wait_visibility() self.grab_set() self.focus_set() self.after(200, lambda: set_window_icon(self)) - - def main_frame(self, master): - # Optimizer - components.label(master, 0, 0, "Optimizer", - tooltip="The type of optimizer") - - # Create the optimizer dropdown menu and set the command - components.options(master, 0, 1, [str(x) for x in list(Optimizer)], self.optimizer_ui_state, "optimizer", - command=self.on_optimizer_change) - - # Defaults Button - components.label(master, 0, 3, "Optimizer Defaults", - tooltip="Load default settings for the selected optimizer") - components.button(self.frame, 0, 4, "Load Defaults", self.load_defaults, - tooltip="Load default settings for the selected optimizer") - - self.create_dynamic_ui(master) - - def clear_dynamic_ui(self, master): + def _rebuild_dynamic_ui(self): with contextlib.suppress(TclError): - for widget in master.winfo_children(): + for widget in self.frame.winfo_children(): grid_info = widget.grid_info() if int(grid_info["row"]) >= 1: widget.destroy() - def create_dynamic_ui( - self, - master, - ): - - # Lookup for the title and tooltip for a key - # @formatter:off - KEY_DETAIL_MAP = { - 'adam_w_mode': {'title': 'Adam W Mode', 'tooltip': 'Whether to use weight decay correction for Adam optimizer.', 'type': 'bool'}, - 'alpha': {'title': 'Alpha', 'tooltip': 'Smoothing parameter for RMSprop and others.', 'type': 'float'}, - 'amsgrad': {'title': 'AMSGrad', 'tooltip': 'Whether to use the AMSGrad variant for Adam.', 'type': 'bool'}, - 'beta1': {'title': 'Beta1', 'tooltip': 'optimizer_momentum term.', 'type': 'float'}, - 'beta2': {'title': 'Beta2', 'tooltip': 'Coefficients for computing running averages of gradient.', 'type': 'float'}, - 'beta3': {'title': 'Beta3', 'tooltip': 'Coefficient for computing the Prodigy stepsize.', 'type': 'float'}, - 'bias_correction': {'title': 'Bias Correction', 'tooltip': 'Whether to use bias correction in optimization algorithms like Adam.', 'type': 'bool'}, - 'block_wise': {'title': 'Block Wise', 'tooltip': 'Whether to perform block-wise model update.', 'type': 'bool'}, - 'capturable': {'title': 'Capturable', 'tooltip': 'Whether some property of the optimizer can be captured.', 'type': 'bool'}, - 'centered': {'title': 'Centered', 'tooltip': 'Whether to center the gradient before scaling. Great for stabilizing the training process.', 'type': 'bool'}, - 'clip_threshold': {'title': 'Clip Threshold', 'tooltip': 'Clipping value for gradients.', 'type': 'float'}, - 'd0': {'title': 'Initial D', 'tooltip': 'Initial D estimate for D-adaptation.', 'type': 'float'}, - 'd_coef': {'title': 'D Coefficient', 'tooltip': 'Coefficient in the expression for the estimate of d.', 'type': 'float'}, - 'dampening': {'title': 'Dampening', 'tooltip': 'Dampening for optimizer_momentum.', 'type': 'float'}, - 'decay_rate': {'title': 'Decay Rate', 'tooltip': 'Rate of decay for moment estimation.', 'type': 'float'}, - 'decouple': {'title': 'Decouple', 'tooltip': 'Use AdamW style optimizer_decoupled weight decay.', 'type': 'bool'}, - 'differentiable': {'title': 'Differentiable', 'tooltip': 'Whether the optimization function is optimizer_differentiable.', 'type': 'bool'}, - 'eps': {'title': 'EPS', 'tooltip': 'A small value to prevent division by zero.', 'type': 'float'}, - 'eps2': {'title': 'EPS 2', 'tooltip': 'A small value to prevent division by zero.', 'type': 'float'}, - 'foreach': {'title': 'ForEach', 'tooltip': 'Whether to use a foreach implementation if available. This implementation is usually faster.', 'type': 'bool'}, - 'fsdp_in_use': {'title': 'FSDP in Use', 'tooltip': 'Flag for using sharded parameters.', 'type': 'bool'}, - 'fused': {'title': 'Fused', 'tooltip': 'Whether to use a fused implementation if available. This implementation is usually faster and requires less memory.', 'type': 'bool'}, - 'fused_back_pass': {'title': 'Fused Back Pass', 'tooltip': 'Whether to fuse the back propagation pass with the optimizer step. This reduces VRAM usage, but is not compatible with gradient accumulation.', 'type': 'bool'}, - 'growth_rate': {'title': 'Growth Rate', 'tooltip': 'Limit for D estimate growth rate.', 'type': 'float'}, - 'initial_accumulator_value': {'title': 'Initial Accumulator Value', 'tooltip': 'Initial value for Adagrad optimizer.', 'type': 'float'}, - 'initial_accumulator': {'title': 'Initial Accumulator', 'tooltip': 'Sets the starting value for both moment estimates to ensure numerical stability and balanced adaptive updates early in training.', 'type': 'float'}, - 'is_paged': {'title': 'Is Paged', 'tooltip': 'Whether the optimizer\'s internal state should be paged to CPU.', 'type': 'bool'}, - 'log_every': {'title': 'Log Every', 'tooltip': 'Intervals at which logging should occur.', 'type': 'int'}, - 'lr_decay': {'title': 'LR Decay', 'tooltip': 'Rate at which learning rate decreases.', 'type': 'float'}, - 'max_unorm': {'title': 'Max Unorm', 'tooltip': 'Maximum value for gradient clipping by norms.', 'type': 'float'}, - 'maximize': {'title': 'Maximize', 'tooltip': 'Whether to optimizer_maximize the optimization function.', 'type': 'bool'}, - 'min_8bit_size': {'title': 'Min 8bit Size', 'tooltip': 'Minimum tensor size for 8-bit quantization.', 'type': 'int'}, - 'quant_block_size': {'title': 'Quant Block Size', 'tooltip': 'Size of a block of normalized 8-bit quantization data. Larger values increase memory efficiency at the cost of data precision.', 'type': 'int'}, - 'momentum': {'title': 'optimizer_momentum', 'tooltip': 'Factor to accelerate SGD in relevant direction.', 'type': 'float'}, - 'nesterov': {'title': 'Nesterov', 'tooltip': 'Whether to enable Nesterov optimizer_momentum.', 'type': 'bool'}, - 'no_prox': {'title': 'No Prox', 'tooltip': 'Whether to use proximity updates or not.', 'type': 'bool'}, - 'optim_bits': {'title': 'Optim Bits', 'tooltip': 'Number of bits used for optimization.', 'type': 'int'}, - 'percentile_clipping': {'title': 'Percentile Clipping', 'tooltip': 'Gradient clipping based on percentile values.', 'type': 'int'}, - 'relative_step': {'title': 'Relative Step', 'tooltip': 'Whether to use a relative step size.', 'type': 'bool'}, - 'safeguard_warmup': {'title': 'Safeguard Warmup', 'tooltip': 'Avoid issues during warm-up stage.', 'type': 'bool'}, - 'scale_parameter': {'title': 'Scale Parameter', 'tooltip': 'Whether to scale the parameter or not.', 'type': 'bool'}, - 'stochastic_rounding': {'title': 'Stochastic Rounding', 'tooltip': 'Stochastic rounding for weight updates. Improves quality when using bfloat16 weights.', 'type': 'bool'}, - 'use_bias_correction': {'title': 'Bias Correction', 'tooltip': 'Turn on Adam\'s bias correction.', 'type': 'bool'}, - 'use_triton': {'title': 'Use Triton', 'tooltip': 'Whether Triton optimization should be used.', 'type': 'bool'}, - 'warmup_init': {'title': 'Warmup Initialization', 'tooltip': 'Whether to warm-up the optimizer initialization.', 'type': 'bool'}, - 'weight_decay': {'title': 'Weight Decay', 'tooltip': 'Regularization to prevent overfitting.', 'type': 'float'}, - 'weight_lr_power': {'title': 'Weight LR Power', 'tooltip': 'During warmup, the weights in the average will be equal to lr raised to this power. Set to 0 for no weighting.', 'type': 'float'}, - 'decoupled_decay': {'title': 'Decoupled Decay', 'tooltip': 'If set as True, then the optimizer uses decoupled weight decay as in AdamW.', 'type': 'bool'}, - 'fixed_decay': {'title': 'Fixed Decay', 'tooltip': '(When Decoupled Decay is True:) Applies fixed weight decay when True; scales decay with learning rate when False.', 'type': 'bool'}, - 'rectify': {'title': 'Rectify', 'tooltip': 'Perform the rectified update similar to RAdam.', 'type': 'bool'}, - 'degenerated_to_sgd': {'title': 'Degenerated to SGD', 'tooltip': 'Performs SGD update when gradient variance is high.', 'type': 'bool'}, - 'k': {'title': 'K', 'tooltip': 'Number of vector projected per iteration.', 'type': 'int'}, - 'xi': {'title': 'Xi', 'tooltip': 'Term used in vector projections to avoid division by zero.', 'type': 'float'}, - 'n_sma_threshold': {'title': 'N SMA Threshold', 'tooltip': 'Number of SMA threshold.', 'type': 'int'}, - 'ams_bound': {'title': 'AMS Bound', 'tooltip': 'Whether to use the AMSBound variant.', 'type': 'bool'}, - 'r': {'title': 'R', 'tooltip': 'EMA factor.', 'type': 'float'}, - 'adanorm': {'title': 'AdaNorm', 'tooltip': 'Whether to use the AdaNorm variant', 'type': 'bool'}, - 'adam_debias': {'title': 'Adam Debias', 'tooltip': 'Only correct the denominator to avoid inflating step sizes early in training.', 'type': 'bool'}, - 'slice_p': {'title': 'Slice parameters', 'tooltip': 'Reduce memory usage by calculating LR adaptation statistics on only every pth entry of each tensor. For values greater than 1 this is an approximation to standard Prodigy. Values ~11 are reasonable.', 'type': 'int'}, - 'cautious': {'title': 'Cautious', 'tooltip': 'Whether to use the Cautious variant', 'type': 'bool'}, - 'weight_decay_by_lr': {'title': 'weight_decay_by_lr', 'tooltip': 'Automatically adjust weight decay based on lr', 'type': 'bool'}, - 'prodigy_steps': {'title': 'prodigy_steps', 'tooltip': 'Turn off Prodigy after N steps', 'type': 'int'}, - 'use_speed': {'title': 'use_speed', 'tooltip': 'use_speed method', 'type': 'bool'}, - 'split_groups': {'title': 'split_groups', 'tooltip': 'Use split groups when training multiple params(uNet,TE..)', 'type': 'bool'}, - 'split_groups_mean': {'title': 'split_groups_mean', 'tooltip': 'Use mean for split groups', 'type': 'bool'}, - 'factored': {'title': 'factored', 'tooltip': 'Use factored', 'type': 'bool'}, - 'factored_fp32': {'title': 'factored_fp32', 'tooltip': 'Use factored_fp32', 'type': 'bool'}, - 'use_stableadamw': {'title': 'use_stableadamw', 'tooltip': 'Use use_stableadamw for gradient scaling', 'type': 'bool'}, - 'use_cautious': {'title': 'use_cautious', 'tooltip': 'Use cautious method', 'type': 'bool'}, - 'use_grams': {'title': 'use_grams', 'tooltip': 'Use grams method', 'type': 'bool'}, - 'use_adopt': {'title': 'use_adopt', 'tooltip': 'Use adopt method', 'type': 'bool'}, - 'd_limiter': {'title': 'd_limiter', 'tooltip': 'Prevent over-estimated LRs when gradients and EMA are still stabilizing', 'type': 'bool'}, - 'use_schedulefree': {'title': 'use_schedulefree', 'tooltip': 'Use Schedulefree method', 'type': 'bool'}, - 'use_orthograd': {'title': 'use_orthograd', 'tooltip': 'Use orthograd method', 'type': 'bool'}, - 'nnmf_factor': {'title': 'Factored Optimizer', 'tooltip': 'Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. This is an approximation which may slightly alter training dynamics.', 'type': 'bool'}, - 'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization.', 'type': 'bool'}, - 'use_atan2': {'title': 'Atan2 Scaling', 'tooltip': 'A robust replacement for eps, which also incorporates gradient clipping, bounding and stabilizing the optimizer updates.', 'type': 'bool'}, - 'use_AdEMAMix': {'title': 'AdEMAMix EMA', 'tooltip': 'Adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, and accelerate the training.', 'type': 'bool'}, - 'beta3_ema': {'title': 'Beta3 EMA', 'tooltip': 'Coefficient for slow-moving EMA of AdEMAMix.', 'type': 'float'}, - 'beta1_warmup': {'title': 'Beta1 Warmup Steps', 'tooltip': 'Number of warmup steps to gradually increase beta1 from Minimum Beta1 Value to its final value. During warmup, beta1 increases linearly. leave it empty to disable warmup and use constant beta1.', 'type': 'int'}, - 'min_beta1': {'title': 'Minimum Beta1', 'tooltip': 'Starting beta1 value for warmup scheduling. Used only when beta1 warmup is enabled. Lower values allow faster initial adaptation, while higher values provide more smoothing. The final beta1 value is specified in the beta1 parameter.', 'type': 'float'}, - 'Simplified_AdEMAMix': {'title': 'Simplified AdEMAMix', 'tooltip': "Enables a simplified, single-EMA variant of AdEMAMix. Instead of blending two moving averages (fast and slow momentum), this version combines the raw current gradient (controlled by 'Grad α') directly with a single theory-based momentum. This makes the optimizer highly responsive to recent gradient information, which can accelerate training in all batch size scenarios when tuned correctly.", 'type': 'bool'}, - 'alpha_grad': {'title': 'Grad α', 'tooltip': 'Controls the mixing coefficient between raw gradients and momentum gradients in Simplified AdEMAMix. Higher values (e.g., 10-100) emphasize recent gradients, suitable for small batch sizes to reduce noise. Lower values (e.g., 0-1) emphasize historical gradients, suitable for large batch sizes for stability. Setting to 0 uses only momentum gradients without raw gradient contribution.', 'type': 'float'}, - 'kourkoutas_beta': {'title': 'Kourkoutas Beta', 'tooltip': 'Enables a layer-wise dynamic β₂ adaptation. This feature makes the optimizer more responsive to "spiky" gradients by lowering β₂ during periods of high variance, and more stable during calm periods by raising β₂ towards its maximum. It can significantly improve training stability and final loss.', 'type': 'bool'}, - 'schedulefree_c': {'title': 'Schedule free averaging strength', 'tooltip': 'Larger values = more responsive (shorter averaging window); smaller values = smoother (longer window). Set to 0 to disable and use the original Schedule-Free rule. Short small batches (≈6-12); long/large-batch (≈50-200).', 'type': 'float'}, - 'ns_steps': {'title': 'Newton-Schulz Iterations', 'tooltip': 'Controls the number of iterations for update orthogonalization. Higher values improve the updates quality but make each step slower. Lower values are faster per step but may be less effective.', 'type': 'int'}, - 'MuonWithAuxAdam': {'title': 'MuonWithAuxAdam', 'tooltip': 'Whether to use the standard way of Muon. Non-hidden layers fallback to ADAMW, and MUON takes the rest. Note: The auxiliary Adam (ADAMW) is typically only relevant for training "full" LoRA (LoRA for all layers) or full finetune and is irrelevant for most common LoRA use cases.', 'type': 'bool'}, - 'muon_hidden_layers': {'title': 'Hidden Layers', 'tooltip': 'Comma-separated list of hidden layers to train using Muon. Regular expressions (if toggled) are supported. Any model layer with a matching name will be trained using Muon. If None is provided it will default to using automatic way of finding hidden layers.', 'type': 'str'}, - 'muon_adam_regex': {'title': 'Use Regex', 'tooltip': 'Whether to use regular expressions for hidden layers.', 'type': 'bool'}, - 'muon_adam_lr': {'title': 'Auxiliary Adam LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer. If empty, it will use the main learning rate.', 'type': 'float'}, - 'muon_te1_adam_lr': {'title': 'AuxAdam TE1 LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer for the first text encoder. If empty, it will use the Auxiliary Adam LR.', 'type': 'float'}, - 'muon_te2_adam_lr': {'title': 'AuxAdam TE2 LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer for the second text encoder. If empty, it will use the Auxiliary Adam LR.', 'type': 'float'}, - 'rms_rescaling': {'title': 'RMS Rescaling', 'tooltip': 'Muon already scales its updates to approximate and use the same learning rate (LR) as Adam. This option integrates a more accurate method to match the Adam LR, but it is slower.', 'type': 'bool'}, - 'normuon_variant': {'title': 'NorMuon Variant', 'tooltip': 'Enables the NorMuon optimizer variant, which combines Muon orthogonalization with per-neuron adaptive learning rates for better convergence and balanced parameter updates. Costs only one scalar state buffer per parameter group, size few KBs, maintaining high memory efficiency.', 'type': 'bool'}, - 'beta2_normuon': {'title': 'NorMuon Beta2', 'tooltip': 'Exponential decay rate for the neuron-wise second-moment estimator in NorMuon (analogous to Adams beta2). Controls how past squared updates influence current normalization.', 'type': 'float'}, - 'low_rank_ortho': {'title': 'Low-rank Orthogonalization', 'tooltip': 'Use low-rank orthogonalization to accelerate Muon by orthogonalizing only in a low-dimensional subspace, improving speed and noise robustness.', 'type': 'bool'}, - 'ortho_rank': {'title': 'Ortho Rank', 'tooltip': 'Target rank for low-rank orthogonalization. Controls the dimensionality of the subspace used for efficient and noise-robust orthogonalization.', 'type': 'int'}, - 'accelerated_ns': {'title': 'Accelerated Newton-Schulz', 'tooltip': 'Applies an enhanced Newton-Schulz variant that replaces heuristic coefficients with optimal coefficients derived at each step. This improves performance and convergence by reducing the number of required operations.', 'type': 'bool'}, - 'cautious_wd': {'title': 'Cautious Weight Decay', 'tooltip': 'Applies weight decay only to parameter coordinates whose signs align with the optimizer update direction. This preserves the original optimization objective while still benefiting from regularization effects, leading to improved convergence and better final performance.', 'type': 'bool'}, - 'approx_mars': {'title': 'Approx MARS-M', 'tooltip': 'Enables Approximated MARS-M, a variance reduction technique. It uses the previous step\'s gradient to correct the current update, leading to lower losses and improved convergence stability. This requires additional state to store the previous gradient.', 'type': 'bool'}, - 'auto_kappa_p': {'title': 'Auto Lion-K', 'tooltip': 'Automatically determines the optimal P-value based on layer dimensions. Uses p=2.0 (Spherical) for 4D (Conv) tensors for stability and rotational invariance, and p=1.0 (Sign) for 2D (Linear) tensors for sparsity. Overrides the manual P-value. Recommend for unet models.', 'type': 'bool'}, - 'compile': {'title': 'Compiled Optimizer', 'tooltip': 'Enables PyTorch compilation for the optimizer internal step logic. This is intended to improve performance by allowing PyTorch to fuse operations and optimize the computational graph.', 'type': 'bool'}, - } - # @formatter:on - - if not self.winfo_exists(): # check if this window isn't open + if not self.winfo_exists(): return - selected_optimizer = self.train_config.optimizer.optimizer - - # Extract the keys for the selected optimizer - for index, key in enumerate(OPTIMIZER_DEFAULT_PARAMETERS[selected_optimizer].keys()): - if key not in KEY_DETAIL_MAP: - continue - arg_info = KEY_DETAIL_MAP[key] - - title = arg_info['title'] - tooltip = arg_info['tooltip'] - type = arg_info['type'] - - row = (index // 2) + 1 - col = 3 * (index % 2) - - components.label(master, row, col, title, tooltip=tooltip) - - if key == 'MuonWithAuxAdam': - frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=row, column=col + 1, columnspan=2, sticky="ew", padx=0, pady=0) - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - - components.switch(frame, 0, 0, self.optimizer_ui_state, key, command=self.update_user_pref) - - self.muon_adam_button = components.button( - frame, 0, 1, "...", self.open_muon_adam_window, - tooltip="Configure the auxiliary AdamW_adv optimizer", - width=20, padx=5 ) - self.toggle_muon_adam_button() - elif type != 'bool': - components.entry(master, row, col + 1, self.optimizer_ui_state, key, - command=self.update_user_pref) - else: - components.switch(master, row, col + 1, self.optimizer_ui_state, key, - command=self.update_user_pref) + self.build_dynamic_content(self.frame, self.controller, self.optimizer_ui_state, + self.update_user_pref, self.open_muon_adam_window) + self.toggle_muon_adam_button() def update_user_pref(self, *args): - update_optimizer_config(self.train_config) + self.controller.on_close() self.toggle_muon_adam_button() def on_optimizer_change(self, *args): - optimizer_config = change_optimizer(self.train_config) - self.ui_state.get_var("optimizer").update(optimizer_config) - - self.clear_dynamic_ui(self.frame) - self.create_dynamic_ui(self.frame) + self.controller.restore_optimizer_config(self.ui_state) + self._rebuild_dynamic_ui() - def load_defaults(self, *args): - optimizer_config = load_optimizer_defaults(self.train_config) - self.ui_state.get_var("optimizer").update(optimizer_config) + def _load_defaults(self, *args): + self.controller.load_defaults(self.ui_state) def on_window_close(self): self.destroy() @@ -261,28 +84,8 @@ def toggle_muon_adam_button(self): self.muon_adam_button.configure(state="normal" if muon_with_adam else "disabled") def open_muon_adam_window(self): - current_optimizer = self.train_config.optimizer.optimizer - - adam_config = TrainOptimizerConfig.default_values() - current_state = self.train_config.optimizer.muon_adam_config - - if current_optimizer == Optimizer.MUON: - defaults = MUON_AUX_ADAM_DEFAULTS - else: - defaults = OPTIMIZER_DEFAULT_PARAMETERS[Optimizer.ADAMW_ADV] - - if current_state is None: - adam_config.from_dict(defaults) - if current_optimizer != Optimizer.MUON: - adam_config.optimizer = Optimizer.ADAMW_ADV - elif isinstance(current_state, dict): - adam_config.from_dict(current_state) - else: - # Should not happen if TrainConfig defines it as dict, but for safety - adam_config = current_state - - temp_adam_ui_state = UIState(self, adam_config) - window = MuonAdamWindow(self, self.train_config, temp_adam_ui_state, current_optimizer) + adam_config, current_optimizer = self.controller.prepare_muon_adam_config() + temp_adam_ui_state = CtkUIState(self, adam_config) + window = CtkMuonAdamWindowView(self, MuonAdamWindowController(self.controller.config, current_optimizer), temp_adam_ui_state) self.wait_window(window) - - self.train_config.optimizer.muon_adam_config = adam_config.to_dict() + self.controller.save_muon_adam_config(adam_config) diff --git a/modules/ui/CtkProfilingWindowView.py b/modules/ui/CtkProfilingWindowView.py index 8d298abe3..15d5055a0 100644 --- a/modules/ui/CtkProfilingWindowView.py +++ b/modules/ui/CtkProfilingWindowView.py @@ -1,16 +1,18 @@ -import faulthandler -from modules.util.ui import components +from modules.ui.BaseProfilingWindowView import BaseProfilingWindowView +from modules.ui.ProfilingWindowController import ProfilingWindowController +from modules.util.ui import ctk_components from modules.util.ui.ui_utils import set_window_icon import customtkinter as ctk -from scalene import scalene_profiler -class ProfilingWindow(ctk.CTkToplevel): - def __init__(self, parent, *args, **kwargs): - super().__init__(parent, *args, **kwargs) - self.parent = parent +class CtkProfilingWindowView(BaseProfilingWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: ProfilingWindowController, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseProfilingWindowView.__init__(self, ctk_components) + + self._controller = controller self.title("Profiling") self.geometry("512x512") @@ -23,35 +25,25 @@ def __init__(self, parent, *args, **kwargs): self.grid_rowconfigure(2, weight=1) self.grid_columnconfigure(0, weight=1) - components.button(self, 0, 0, "Dump stack", self._dump_stack) - self._profile_button = components.button( - self, 1, 0, "Start Profiling", self._start_profiler, - tooltip="Turns on/off Scalene profiling. Only works when OneTrainer is launched with Scalene!") - # Bottom bar self._bottom_bar = ctk.CTkFrame(master=self, corner_radius=0) self._bottom_bar.grid(row=2, column=0, sticky="sew") - self._message_label = components.label(self._bottom_bar, 0, 0, "Inactive") + + self.build_content(self, self._bottom_bar, controller) self.protocol("WM_DELETE_WINDOW", self.withdraw) self.withdraw() self.after(200, lambda: set_window_icon(self)) - def _dump_stack(self): - with open('stacks.txt', 'w') as f: - faulthandler.dump_traceback(f) - self._message_label.configure(text='Stack dumped to stacks.txt') - - def _end_profiler(self): - scalene_profiler.stop() - - self._message_label.configure(text='Inactive') - self._profile_button.configure(text='Start Profiling') - self._profile_button.configure(command=self._start_profiler) - - def _start_profiler(self): - scalene_profiler.start() - - self._message_label.configure(text='Profiling active...') - self._profile_button.configure(text='End Profiling') - self._profile_button.configure(command=self._end_profiler) + def set_message(self, text): + self._message_label.configure(text=text) + + def set_profiling_active(self, active): + if active: + self._message_label.configure(text='Profiling active...') + self._profile_button.configure(text='End Profiling') + self._profile_button.configure(command=self._controller.end_profiler) + else: + self._message_label.configure(text='Inactive') + self._profile_button.configure(text='Start Profiling') + self._profile_button.configure(command=self._controller.start_profiler) diff --git a/modules/ui/CtkSampleFrameView.py b/modules/ui/CtkSampleFrameView.py index 297caac29..167b25692 100644 --- a/modules/ui/CtkSampleFrameView.py +++ b/modules/ui/CtkSampleFrameView.py @@ -1,37 +1,28 @@ -from modules.util.config.SampleConfig import SampleConfig -from modules.util.enum.ModelType import ModelType -from modules.util.enum.NoiseScheduler import NoiseScheduler -from modules.util.ui import components -from modules.util.ui.UIState import UIState +from modules.ui.BaseSampleFrameView import BaseSampleFrameView +from modules.ui.SampleFrameController import SampleFrameController +from modules.util.ui import ctk_components import customtkinter as ctk -class SampleFrame(ctk.CTkFrame): +class CtkSampleFrameView(BaseSampleFrameView, ctk.CTkFrame): def __init__( self, parent, - sample: SampleConfig, - ui_state: UIState, - model_type: ModelType, + controller: SampleFrameController, + ui_state, include_prompt: bool = True, include_settings: bool = True, ): ctk.CTkFrame.__init__(self, parent, fg_color="transparent") - - self.sample = sample - self.ui_state = ui_state - self.model_type = model_type - - is_flow_matching = model_type.is_flow_matching() - is_inpainting_model = model_type.has_conditioning_image_input() - is_video_model = model_type.is_video_model() + BaseSampleFrameView.__init__(self, ctk_components) if include_prompt and include_prompt: self.grid_rowconfigure(0, weight=0) self.grid_rowconfigure(1, weight=1) self.grid_columnconfigure(0, weight=1) + top_frame = None if include_prompt: top_frame = ctk.CTkFrame(self, fg_color="transparent") top_frame.grid(row=0, column=0, padx=0, pady=0, sticky="nsew") @@ -39,6 +30,7 @@ def __init__( top_frame.grid_columnconfigure(0, weight=0) top_frame.grid_columnconfigure(1, weight=1) + bottom_frame = None if include_settings: bottom_frame = ctk.CTkFrame(self, fg_color="transparent") bottom_frame.grid(row=1, column=0, padx=0, pady=0, sticky="nsew") @@ -48,87 +40,4 @@ def __init__( bottom_frame.grid_columnconfigure(2, weight=0) bottom_frame.grid_columnconfigure(3, weight=1) - if include_prompt: - # prompt - components.label(top_frame, 0, 0, "prompt:") - components.entry(top_frame, 0, 1, self.ui_state, "prompt") - - # negative prompt - components.label(top_frame, 1, 0, "negative prompt:") - components.entry(top_frame, 1, 1, self.ui_state, "negative_prompt") - - if include_settings: - # width - components.label(bottom_frame, 0, 0, "width:") - components.entry(bottom_frame, 0, 1, self.ui_state, "width") - - # height - components.label(bottom_frame, 0, 2, "height:") - components.entry(bottom_frame, 0, 3, self.ui_state, "height") - - if is_video_model: - # frames - components.label(bottom_frame, 1, 0, "frames:", - tooltip="Number of frames to generate. Only used when generating videos.") - components.entry(bottom_frame, 1, 1, self.ui_state, "frames") - - # length - components.label(bottom_frame, 1, 2, "length:", - tooltip="Length in seconds of audio output.") - components.entry(bottom_frame, 1, 3, self.ui_state, "length") - - # seed - components.label(bottom_frame, 2, 0, "seed:") - components.entry(bottom_frame, 2, 1, self.ui_state, "seed") - - # random seed - components.label(bottom_frame, 2, 2, "random seed:") - components.switch(bottom_frame, 2, 3, self.ui_state, "random_seed") - - # cfg scale - components.label(bottom_frame, 3, 0, "cfg scale:") - components.entry(bottom_frame, 3, 1, self.ui_state, "cfg_scale") - - # sampler - if not is_flow_matching: - components.label(bottom_frame, 4, 2, "sampler:") - components.options_kv(bottom_frame, 4, 3, [ - ("DDIM", NoiseScheduler.DDIM), - ("Euler", NoiseScheduler.EULER), - ("Euler A", NoiseScheduler.EULER_A), - # ("DPM++", NoiseScheduler.DPMPP), # TODO: produces noisy samples - # ("DPM++ SDE", NoiseScheduler.DPMPP_SDE), # TODO: produces noisy samples - ("UniPC", NoiseScheduler.UNIPC), - ("Euler Karras", NoiseScheduler.EULER_KARRAS), - ("DPM++ Karras", NoiseScheduler.DPMPP_KARRAS), - ("DPM++ SDE Karras", NoiseScheduler.DPMPP_SDE_KARRAS), - ("UniPC Karras", NoiseScheduler.UNIPC_KARRAS) - ], self.ui_state, "noise_scheduler") - - # steps - components.label(bottom_frame, 4, 0, "steps:") - components.entry(bottom_frame, 4, 1, self.ui_state, "diffusion_steps") - - # inpainting - if is_inpainting_model: - components.label(bottom_frame, 5, 0, "inpainting:", - tooltip="Enables inpainting sampling. Only available when sampling from an inpainting model.") - components.switch(bottom_frame, 5, 1, self.ui_state, "sample_inpainting") - - # 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", - mode="file", - allow_model_files=False, - allow_image_files=True, - ) - - # 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", - mode="file", - allow_model_files=False, - allow_image_files=True, - ) + self.build_content(top_frame, bottom_frame, ui_state, controller, include_prompt, include_settings) diff --git a/modules/ui/CtkSampleParamsWindowView.py b/modules/ui/CtkSampleParamsWindowView.py index 2b0b3f3f1..8229a19f6 100644 --- a/modules/ui/CtkSampleParamsWindowView.py +++ b/modules/ui/CtkSampleParamsWindowView.py @@ -1,20 +1,17 @@ -from modules.ui.SampleFrame import SampleFrame -from modules.util.config.SampleConfig import SampleConfig -from modules.util.enum.ModelType import ModelType -from modules.util.ui import components +from modules.ui.BaseSampleParamsWindowView import BaseSampleParamsWindowView +from modules.ui.CtkSampleFrameView import CtkSampleFrameView +from modules.ui.SampleFrameController import SampleFrameController +from modules.ui.SampleParamsWindowController import SampleParamsWindowController +from modules.util.ui import ctk_components from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState import customtkinter as ctk -class SampleParamsWindow(ctk.CTkToplevel): - def __init__(self, parent, sample: SampleConfig, ui_state: UIState, model_type: ModelType | None = None, *args, **kwargs): - super().__init__(parent, *args, **kwargs) - - self.sample = sample - self.ui_state = ui_state - self.model_type = model_type +class CtkSampleParamsWindowView(BaseSampleParamsWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: SampleParamsWindowController, ui_state, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseSampleParamsWindowView.__init__(self, ctk_components) self.title("Sample") self.geometry("800x500") @@ -24,16 +21,12 @@ def __init__(self, parent, sample: SampleConfig, ui_state: UIState, model_type: self.grid_rowconfigure(1, weight=0) self.grid_columnconfigure(0, weight=1) - frame = SampleFrame(self, self.sample, self.ui_state, model_type=model_type) + frame = CtkSampleFrameView(self, SampleFrameController(controller.sample, controller.model_type), ui_state) frame.grid(row=0, column=0, padx=0, pady=0, sticky="nsew") - components.button(self, 1, 0, "ok", self.__ok) + self.components.button(self, 1, 0, "ok", self.destroy) self.wait_visibility() self.grab_set() self.focus_set() self.after(200, lambda: set_window_icon(self)) - - - def __ok(self): - self.destroy() diff --git a/modules/ui/CtkSampleWindowView.py b/modules/ui/CtkSampleWindowView.py index 0f91ad2fa..3b67f03d5 100644 --- a/modules/ui/CtkSampleWindowView.py +++ b/modules/ui/CtkSampleWindowView.py @@ -1,75 +1,43 @@ import contextlib -import copy -import os import tkinter as tk import traceback -from modules.model.BaseModel import BaseModel from modules.modelSampler.BaseModelSampler import ( - BaseModelSampler, ModelSamplerOutput, ) -from modules.ui.SampleFrame import SampleFrame -from modules.util import create -from modules.util.callbacks.TrainCallbacks import TrainCallbacks -from modules.util.commands.TrainCommands import TrainCommands -from modules.util.config.SampleConfig import SampleConfig -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.EMAMode import EMAMode +from modules.ui.BaseSampleWindowView import BaseSampleWindowView +from modules.ui.CtkSampleFrameView import CtkSampleFrameView +from modules.ui.SampleFrameController import SampleFrameController +from modules.ui.SampleWindowController import SampleWindowController from modules.util.enum.FileType import FileType -from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.time_util import get_string_timestamp -from modules.util.ui import components +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState - -import torch import customtkinter as ctk from PIL import Image -class SampleWindow(ctk.CTkToplevel): +class CtkSampleWindowView(BaseSampleWindowView, ctk.CTkToplevel): def __init__( self, parent, - train_config: TrainConfig, - use_external_model: bool, - callbacks: TrainCallbacks | None = None, - commands: TrainCommands | None = None, + controller: SampleWindowController, *args, **kwargs ): - super().__init__(parent, *args, **kwargs) + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseSampleWindowView.__init__(self, ctk_components) self.title("Sample") self.geometry("1200x800") self.resizable(True, True) - if not use_external_model: - self.initial_train_config = TrainConfig.default_values().from_dict(train_config.to_dict()) - # remove some settings to speed up model loading for sampling - self.initial_train_config.optimizer.optimizer = None - self.initial_train_config.ema = EMAMode.OFF - else: - self.initial_train_config = None - - #TODO why is there a current_train_config and an initial_train_config? - #current_train_config doesn't seem to ever change - self.current_train_config = train_config - self.callbacks = callbacks - self.commands = commands - - # get model specific defaults - model_type = train_config.model_type - self.sample = SampleConfig.default_values(model_type) - self.ui_state = UIState(self, self.sample) - - if use_external_model: - self.callbacks.set_on_sample_custom(self.__update_preview) - self.callbacks.set_on_update_sample_custom_progress(self.__update_progress) - else: - self.model = None - self.model_sampler = None + model_type = controller.get_model_type() + self.ui_state = CtkUIState(self, controller.sample) + + if controller.use_external_model: + controller.callbacks.set_on_sample_custom(self.__update_preview) + controller.callbacks.set_on_update_sample_custom_progress(self.__update_progress) self.grid_rowconfigure(0, weight=0) self.grid_rowconfigure(1, weight=1) @@ -78,10 +46,10 @@ def __init__( self.grid_columnconfigure(0, weight=0) self.grid_columnconfigure(1, weight=1) - prompt_frame = SampleFrame(self, self.sample, self.ui_state, include_settings=False, model_type=model_type) + prompt_frame = CtkSampleFrameView(self, SampleFrameController(controller.sample, model_type), self.ui_state, include_settings=False) prompt_frame.grid(row=0, column=0, columnspan=2, padx=0, pady=0, sticky="nsew") - settings_frame = SampleFrame(self, self.sample, self.ui_state, include_prompt=False, model_type=model_type) + settings_frame = CtkSampleFrameView(self, SampleFrameController(controller.sample, model_type), self.ui_state, include_prompt=False) settings_frame.grid(row=1, column=0, padx=0, pady=0, sticky="nsew") # image @@ -93,70 +61,14 @@ def __init__( image_label = ctk.CTkLabel(master=self, text="", image=self.image, height=512, width=512) image_label.grid(row=1, column=1, rowspan=3, sticky="nsew") - self.progress = components.progress(self, 2, 0) - components.button(self, 3, 0, "sample", self.__sample) + self.progress = self.components.progress(self, 2, 0) + self.components.button(self, 3, 0, "sample", + lambda: controller.do_sample(self.__update_preview, self.__update_progress)) self.wait_visibility() self.focus_set() self.after(200, lambda: set_window_icon(self)) - def __load_model(self) -> BaseModel: - model_loader = create.create_model_loader( - model_type=self.initial_train_config.model_type, - training_method=self.initial_train_config.training_method, - ) - - model_setup = create.create_model_setup( - model_type=self.initial_train_config.model_type, - train_device=torch.device(self.initial_train_config.train_device), - temp_device=torch.device(self.initial_train_config.temp_device), - training_method=self.initial_train_config.training_method, - ) - - model_names = self.initial_train_config.model_names() - if self.initial_train_config.continue_last_backup: - last_backup_path = self.initial_train_config.get_last_backup_path() - - if last_backup_path: - if self.initial_train_config.training_method == TrainingMethod.LORA: - model_names.lora = last_backup_path - elif self.initial_train_config.training_method == TrainingMethod.EMBEDDING: - model_names.embedding.model_name = last_backup_path - else: # fine-tunes - model_names.base_model = last_backup_path - - print(f"Loading from backup '{last_backup_path}'...") - else: - print("No backup found, loading without backup...") - - if self.initial_train_config.quantization.cache_dir is None: - self.initial_train_config.quantization.cache_dir = self.initial_train_config.cache_dir + "/quantization" - os.makedirs(self.initial_train_config.quantization.cache_dir, exist_ok=True) - - model = model_loader.load( - model_type=self.initial_train_config.model_type, - model_names=model_names, - weight_dtypes=self.initial_train_config.weight_dtypes(), - quantization=self.initial_train_config.quantization, - ) - model.train_config = self.initial_train_config - - model_setup.setup_optimizations(model, self.initial_train_config) - model_setup.setup_train_device(model, self.initial_train_config) - model_setup.setup_model(model, self.initial_train_config) - model.to(torch.device(self.initial_train_config.temp_device)) - - return model - - def __create_sampler(self, model: BaseModel) -> BaseModelSampler: - return create.create_model_sampler( - train_device=torch.device(self.initial_train_config.train_device), - temp_device=torch.device(self.initial_train_config.temp_device), - model=model, - model_type=self.initial_train_config.model_type, - training_method=self.initial_train_config.training_method, - ) - def __update_preview(self, sampler_output: ModelSamplerOutput): if sampler_output.file_type == FileType.IMAGE: image = sampler_output.data @@ -172,43 +84,6 @@ def __update_progress(self, progress: int, max_progress: int): def __dummy_image(self) -> Image: return Image.new(mode="RGB", size=(512, 512), color=(0, 0, 0)) - def __sample(self): - sample = copy.copy(self.sample) - - if self.commands: - self.commands.sample_custom(sample) - else: - if self.model is None: - # lazy initialization - self.model = self.__load_model() - self.model_sampler = self.__create_sampler(self.model) - - sample.from_train_config(self.current_train_config) - - sample_dir = os.path.join( - self.initial_train_config.workspace_dir, - "samples", - "custom", - ) - - progress = self.model.train_progress - sample_path = os.path.join( - sample_dir, - f"{get_string_timestamp()}-training-sample-{progress.filename_string()}" - ) - - self.model.eval() - - self.model_sampler.sample( - sample_config=sample, - destination=sample_path, - image_format=self.current_train_config.sample_image_format, - video_format=self.current_train_config.sample_video_format, - audio_format=self.current_train_config.sample_audio_format, - on_sample=self.__update_preview, - on_update_progress=self.__update_progress, - ) - def destroy(self): try: if hasattr(self, "_icon_image_ref"): diff --git a/modules/ui/CtkSamplingTabView.py b/modules/ui/CtkSamplingTabView.py index 5a3c44f08..dfe1a704a 100644 --- a/modules/ui/CtkSamplingTabView.py +++ b/modules/ui/CtkSamplingTabView.py @@ -1,20 +1,17 @@ -from modules.ui.ConfigList import ConfigList -from modules.ui.SampleParamsWindow import SampleParamsWindow -from modules.util.config.SampleConfig import SampleConfig -from modules.util.config.TrainConfig import TrainConfig -from modules.util.ui import components -from modules.util.ui.UIState import UIState +from modules.ui.BaseSamplingTabView import BaseSampleWidgetView, BaseSamplingTabView +from modules.ui.CtkConfigListView import CtkConfigListView +from modules.ui.CtkSampleParamsWindowView import CtkSampleParamsWindowView +from modules.ui.SamplingTabController import SamplingTabController +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState import customtkinter as ctk -class SamplingTab(ConfigList): - - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__( - master, - train_config, - ui_state, +class CtkSamplingTabView(CtkConfigListView, BaseSamplingTabView): + def __init__(self, master, controller: SamplingTabController, ui_state): + CtkConfigListView.__init__( + self, master, controller, ui_state, from_external_file=True, attr_name="sample_definition_file_name", config_dir="training_samples", @@ -22,103 +19,32 @@ 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, ) - def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - return SampleWidget(master, element, i, open_command, remove_command, clone_command, save_command) - - def create_new_element(self) -> dict: - return SampleConfig.default_values(self.train_config.model_type) - def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: - return SampleParamsWindow(self.master, self.current_config[i], ui_state, model_type=self.train_config.model_type) + return self.controller.open_element_window(self.master, self.current_config[i], ui_state, CtkSampleParamsWindowView) + + def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): + return CtkSampleWidgetView(master, element, i, open_command, remove_command, clone_command, save_command) -class SampleWidget(ctk.CTkFrame): +class CtkSampleWidgetView(BaseSampleWidgetView, ctk.CTkFrame): def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): - super().__init__( - master=master, corner_radius=10, bg_color="transparent" - ) + ctk.CTkFrame.__init__(self, master=master, corner_radius=10, bg_color="transparent") + BaseSampleWidgetView.__init__(self, ctk_components) - self.element = element - self.ui_state = UIState(self, element) - self.i = i - self.save_command = save_command + self.ui_state = CtkUIState(self, element) self.grid_columnconfigure(10, weight=1) - # close button - close_button = ctk.CTkButton( - master=self, - width=20, - height=20, - text="X", - corner_radius=2, - fg_color="#C00000", - command=lambda: remove_command(self.i), - ) - close_button.grid(row=0, column=0) - - # clone button - clone_button = ctk.CTkButton( - master=self, - width=20, - height=20, - text="+", - corner_radius=2, - fg_color="#00C000", - command=lambda: clone_command(self.i), - ) - clone_button.grid(row=0, column=1, padx=5) + self.build_content(self, element, self.ui_state, i, open_command, remove_command, clone_command, save_command) - # enabled - self.enabled_switch = components.switch(self, 0, 2, self.ui_state, "enabled", self.__switch_enabled) - self.enabled_switch.configure(width=40) - - # width - components.label(self, 0, 3, "width:") - self.width_entry = components.entry(self, 0, 4, self.ui_state, "width") + def _bind_save(self, save_command): self.width_entry.bind('', lambda _: save_command()) - self.width_entry.configure(width=50) - - # height - components.label(self, 0, 5, "height:") - self.height_entry = components.entry(self, 0, 6, self.ui_state, "height") self.height_entry.bind('', lambda _: save_command()) - self.height_entry.configure(width=50) - - # seed - components.label(self, 0, 7, "seed:") - self.seed_entry = components.entry(self, 0, 8, self.ui_state, "seed") self.seed_entry.bind('', lambda _: save_command()) - self.seed_entry.configure(width=80) - - # prompt - components.label(self, 0, 9, "prompt:") - self.prompt_entry = components.entry(self, 0, 10, self.ui_state, "prompt") self.prompt_entry.bind('', lambda _: save_command()) - # button - self.button = components.icon_button(self, 0, 11, "...", lambda: open_command(self.i, self.ui_state)) - self.button.configure(width=40) - - self.__set_enabled() - - def __switch_enabled(self): - self.save_command() - self.__set_enabled() - - def __set_enabled(self): - enabled = self.element.enabled - self.width_entry.configure(state="normal" if enabled else "disabled") - self.height_entry.configure(state="normal" if enabled else "disabled") - self.prompt_entry.configure(state="normal" if enabled else "disabled") - self.seed_entry.configure(state="normal" if enabled else "disabled") - self.button.configure(state="normal" if enabled else "disabled") - - def configure_element(self): - pass - def place_in_list(self): self.grid(row=self.i, column=0, pady=5, padx=5, sticky="new") diff --git a/modules/ui/CtkSchedulerParamsWindowView.py b/modules/ui/CtkSchedulerParamsWindowView.py index f96ed4876..9a7c41b96 100644 --- a/modules/ui/CtkSchedulerParamsWindowView.py +++ b/modules/ui/CtkSchedulerParamsWindowView.py @@ -1,24 +1,23 @@ -from modules.ui.ConfigList import ConfigList -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.LearningRateScheduler import LearningRateScheduler -from modules.util.ui import components +from modules.ui.BaseSchedulerParamsWindowView import BaseKvParamsView, BaseSchedulerParamsWindowView +from modules.ui.CtkConfigListView import CtkConfigListView +from modules.ui.SchedulerParamsWindowController import KvParamsController, SchedulerParamsWindowController +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState import customtkinter as ctk -class KvParams(ConfigList): - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__( - master, - train_config, - ui_state, +class CtkKvParamsView(CtkConfigListView, BaseKvParamsView): + def __init__(self, master, controller: KvParamsController, ui_state): + CtkConfigListView.__init__( + self, master, controller, ui_state, attr_name="scheduler_params", from_external_file=False, add_button_text="add parameter", - is_full_width=True + is_full_width=True, ) + BaseKvParamsView.__init__(self, ctk_components) def refresh_ui(self): self._create_element_list() @@ -26,18 +25,12 @@ def refresh_ui(self): def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): return KvWidget(master, element, i, open_command, remove_command, clone_command, save_command) - def create_new_element(self) -> dict[str, str]: - return {"key": "", "value": ""} - - def open_element_window(self, i, ui_state): - pass - class KvWidget(ctk.CTkFrame): def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): super().__init__(master=master, bg_color="transparent") self.element = element - self.ui_state = UIState(self, element) + self.ui_state = CtkUIState(self, element) self.i = i self.save_command = save_command @@ -57,14 +50,14 @@ def __init__(self, master, element, i, open_command, remove_command, clone_comma # Key tooltip_key = "Key name for an argument in your scheduler" - self.key = components.entry(self, 0, 1, self.ui_state, "key", + self.key = ctk_components.entry(self, 0, 1, self.ui_state, "key", tooltip=tooltip_key, wide_tooltip=True) self.key.bind("", lambda _: save_command()) self.key.configure(width=50) # Value tooltip_val = "Value for an argument in your scheduler. Some special values can be used, wrapped in percent signs: LR, EPOCHS, STEPS_PER_EPOCH, TOTAL_STEPS, SCHEDULER_STEPS. Note that OneTrainer calls step() after every individual learning step, not every epoch, so what Torch calls 'epoch' you should treat as 'step'." - self.value = components.entry(self, 0, 2, self.ui_state, "value", + self.value = ctk_components.entry(self, 0, 2, self.ui_state, "value", tooltip=tooltip_val, wide_tooltip=True) self.value.bind("", lambda _: save_command()) self.value.configure(width=50) @@ -73,47 +66,31 @@ def place_in_list(self): self.grid(row=self.i, column=0, padx=5, pady=5, sticky="new") -class SchedulerParamsWindow(ctk.CTkToplevel): - def __init__(self, parent, train_config: TrainConfig, ui_state, *args, **kwargs): - super().__init__(parent, *args, **kwargs) - - self.parent = parent - self.train_config = train_config - self.ui_state = ui_state +class CtkSchedulerParamsWindowView(BaseSchedulerParamsWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: SchedulerParamsWindowController, ui_state, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseSchedulerParamsWindowView.__init__(self, ctk_components) self.title("Learning Rate Scheduler Settings") self.geometry("800x400") self.resizable(True, True) - self.grid_rowconfigure(0, weight=1) self.grid_rowconfigure(1, weight=0) self.grid_columnconfigure(0, weight=1) - self.frame = ctk.CTkFrame(self) - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) + frame = ctk.CTkFrame(self) + frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) - self.expand_frame = ctk.CTkFrame(self.frame, bg_color="transparent") - self.expand_frame.grid(row=1, column=0, columnspan=2, sticky="nsew") + expand_frame = ctk.CTkFrame(frame, bg_color="transparent") + expand_frame.grid(row=1, column=0, columnspan=2, sticky="nsew") - components.button(self, 1, 0, "ok", command=self.on_window_close) - self.main_frame(self.frame) + self.components.button(self, 1, 0, "ok", command=self.destroy) + self.build_content(frame, controller, ui_state) + CtkKvParamsView(expand_frame, KvParamsController(controller.config), ui_state) self.wait_visibility() self.grab_set() self.focus_set() self.after(200, lambda: set_window_icon(self)) - - - def main_frame(self, master): - if self.train_config.learning_rate_scheduler is LearningRateScheduler.CUSTOM: - components.label(master, 0, 0, "Class Name", - tooltip="Python class module and name for the custom scheduler class, in the form of ..") - components.entry(master, 0, 1, self.ui_state, "custom_learning_rate_scheduler") - - # Any additional parameters, in key-value form. - self.params = KvParams(self.expand_frame, self.train_config, self.ui_state) - - def on_window_close(self): - self.destroy() diff --git a/modules/ui/CtkTimestepDistributionWindowView.py b/modules/ui/CtkTimestepDistributionWindowView.py index 21e41ce3e..69f2ae7d3 100644 --- a/modules/ui/CtkTimestepDistributionWindowView.py +++ b/modules/ui/CtkTimestepDistributionWindowView.py @@ -1,15 +1,8 @@ -from modules.modelSetup.mixin.ModelSetupNoiseMixin import ( - ModelSetupNoiseMixin, -) -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.TimestepDistribution import TimestepDistribution -from modules.util.ui import components +from modules.ui.BaseTimestepDistributionWindowView import BaseTimestepDistributionWindowView +from modules.ui.TimestepDistributionWindowController import TimestepDistributionWindowController +from modules.util.ui import ctk_components from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState - -import torch -from torch import Tensor import customtkinter as ctk from customtkinter import AppearanceModeTracker, ThemeManager @@ -17,127 +10,38 @@ from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg -class TimestepGenerator(ModelSetupNoiseMixin): - - def __init__( - self, - timestep_distribution: TimestepDistribution, - min_noising_strength: float, - max_noising_strength: float, - noising_weight: float, - noising_bias: float, - timestep_shift: float, - ): - super().__init__() - - self.timestep_distribution = timestep_distribution - self.min_noising_strength = min_noising_strength - self.max_noising_strength = max_noising_strength - self.noising_weight = noising_weight - self.noising_bias = noising_bias - self.timestep_shift = timestep_shift - - def generate(self) -> Tensor: - generator = torch.Generator() - generator.seed() - - config = TrainConfig.default_values() - config.timestep_distribution = self.timestep_distribution - config.min_noising_strength = self.min_noising_strength - config.max_noising_strength = self.max_noising_strength - config.noising_weight = self.noising_weight - config.noising_bias = self.noising_bias - config.timestep_shift = self.timestep_shift - - - return self._get_timestep_discrete( - num_train_timesteps=1000, - deterministic=False, - generator=generator, - batch_size=1000000, - config=config, - ) - - -class TimestepDistributionWindow(ctk.CTkToplevel): +class CtkTimestepDistributionWindowView(BaseTimestepDistributionWindowView, ctk.CTkToplevel): def __init__( self, parent, - config: TrainConfig, - ui_state: UIState, + controller: TimestepDistributionWindowController, + ui_state, *args, **kwargs, ): - super().__init__(parent, *args, **kwargs) + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseTimestepDistributionWindowView.__init__(self, ctk_components) self.title("Timestep Distribution") self.geometry("900x600") self.resizable(True, True) - self.config = config - self.ui_state = ui_state - self.image_preview_file_index = 0 + self.controller = controller self.ax = None self.canvas = None self.grid_rowconfigure(0, weight=1) self.grid_columnconfigure(0, weight=1) - frame = self.__content_frame(self) - frame.grid(row=0, column=0, sticky='nsew') - components.button(self, 1, 0, "ok", self.__ok) - - self.wait_visibility() - self.after(200, lambda: set_window_icon(self)) - self.grab_set() - self.focus_set() - - def __content_frame(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame = ctk.CTkScrollableFrame(self, fg_color="transparent") frame.grid_columnconfigure(0, weight=0) frame.grid_columnconfigure(1, weight=0) frame.grid_columnconfigure(2, weight=0) frame.grid_columnconfigure(3, weight=1) frame.grid_rowconfigure(7, weight=1) - # timestep distribution - components.label(frame, 0, 0, "Timestep Distribution", - tooltip="Selects the function to sample timesteps during training", - wide_tooltip=True) - components.options(frame, 0, 1, [str(x) for x in list(TimestepDistribution)], self.ui_state, - "timestep_distribution") - - # min noising strength - components.label(frame, 1, 0, "Min Noising Strength", - tooltip="Specifies the minimum noising strength used during training. This can help to improve composition, but prevents finer details from being trained") - components.entry(frame, 1, 1, self.ui_state, "min_noising_strength") - - # max noising strength - components.label(frame, 2, 0, "Max Noising Strength", - tooltip="Specifies the maximum noising strength used during training. This can be useful to reduce overfitting, but also reduces the impact of training samples on the overall image composition") - components.entry(frame, 2, 1, self.ui_state, "max_noising_strength") - - # noising weight - components.label(frame, 3, 0, "Noising Weight", - tooltip="Controls the weight parameter of the timestep distribution function. Use the preview to see more details.") - components.entry(frame, 3, 1, self.ui_state, "noising_weight") - - # noising bias - components.label(frame, 4, 0, "Noising Bias", - tooltip="Controls the bias parameter of the timestep distribution function. Use the preview to see more details.") - components.entry(frame, 4, 1, self.ui_state, "noising_bias") - - # timestep shift - components.label(frame, 5, 0, "Timestep Shift", - tooltip="Shift the timestep distribution. Use the preview to see more details.") - components.entry(frame, 5, 1, self.ui_state, "timestep_shift") - - # dynamic timestep shifting - components.label(frame, 6, 0, "Dynamic Timestep Shifting", - tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Dynamic Timestep Shifting is not shown in the preview. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) - components.switch(frame, 6, 1, self.ui_state, "dynamic_timestep_shifting") - - - # plot + self.build_content(frame, controller, ui_state) + + # matplotlib chart (CTK-only: needs winfo_rgb from the toplevel) appearance_mode = AppearanceModeTracker.get_mode() background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) @@ -163,24 +67,18 @@ def __content_frame(self, master): self.__update_preview() # update button - components.button(frame, 8, 3, "Update Preview", command=self.__update_preview) + ctk_components.button(frame, 8, 3, "Update Preview", command=self.__update_preview) frame.pack(fill="both", expand=1) - return frame + frame.grid(row=0, column=0, sticky='nsew') + ctk_components.button(self, 1, 0, "ok", self.destroy) - def __update_preview(self): - generator = TimestepGenerator( - timestep_distribution=self.config.timestep_distribution, - min_noising_strength=self.config.min_noising_strength, - max_noising_strength=self.config.max_noising_strength, - noising_weight=self.config.noising_weight, - noising_bias=self.config.noising_bias, - timestep_shift=self.config.timestep_shift, - ) + self.wait_visibility() + self.after(200, lambda: set_window_icon(self)) + self.grab_set() + self.focus_set() + def __update_preview(self): self.ax.cla() - self.ax.hist(generator.generate(), bins=1000, range=(0, 999)) + self.ax.hist(self.controller.generate_preview_data(), bins=1000, range=(0, 999)) self.canvas.draw() - - def __ok(self): - self.destroy() diff --git a/modules/ui/CtkTopBarView.py b/modules/ui/CtkTopBarView.py index 820fdb71a..ee1bcfa75 100644 --- a/modules/ui/CtkTopBarView.py +++ b/modules/ui/CtkTopBarView.py @@ -1,260 +1,50 @@ -import json -import os -import traceback -import webbrowser from collections.abc import Callable -from contextlib import suppress -from modules.util import path_util -from modules.util.config.SecretsConfig import SecretsConfig -from modules.util.config.TrainConfig import TrainConfig +from modules.ui.BaseTopBarView import BaseTopBarView +from modules.ui.TopBarController import TopBarController from modules.util.enum.ModelType import ModelType from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.optimizer_util import change_optimizer -from modules.util.path_util import write_json_atomic -from modules.util.ui import components, dialogs -from modules.util.ui.UIState import UIState +from modules.util.ui import ctk_components, dialogs +from modules.util.ui.CtkUIState import CtkUIState import customtkinter as ctk -class TopBar: +class CtkTopBarView(BaseTopBarView): def __init__( self, master, - train_config: TrainConfig, - ui_state: UIState, + controller: TopBarController, + ui_state, change_model_type_callback: Callable[[ModelType], None], change_training_method_callback: Callable[[TrainingMethod], None], load_preset_callback: Callable[[], None], ): - self.master = master - self.train_config = train_config - self.ui_state = ui_state - self.change_model_type_callback = change_model_type_callback - self.change_training_method_callback = change_training_method_callback - self.load_preset_callback = load_preset_callback + BaseTopBarView.__init__(self, ctk_components) - self.dir = "training_presets" + frame = ctk.CTkFrame(master=master, corner_radius=0) + frame.grid(row=0, column=0, sticky="nsew") - self.config_ui_data = { - "config_name": path_util.canonical_join(self.dir, "#.json") - } - self.config_ui_state = UIState(master, self.config_ui_data) + self.build(frame, master, controller, ui_state, change_model_type_callback, change_training_method_callback, load_preset_callback) - self.configs = [("", path_util.canonical_join(self.dir, "#.json"))] - self.__load_available_config_names() + def _make_config_ui_state(self, master, data): + return CtkUIState(master, data) - self.current_config = [] + def _get_dropdown_text(self, widget) -> str: + return widget.get() - self.frame = ctk.CTkFrame(master=master, corner_radius=0) - self.frame.grid(row=0, column=0, sticky="nsew") - - self.training_method = None - - # title - components.app_title(self.frame, 0, 0) - - # dropdown - self.configs_dropdown = None - self.__create_configs_dropdown() - - # remove button - # TODO - # components.icon_button(self.frame, 0, 2, "-", self.__remove_config) - - # Wiki button - components.button(self.frame, 0, 4, "Wiki", self.open_wiki, width=50) - - # save button - components.button(self.frame, 0, 3, "Save config", self.__save_config, - tooltip="Save the current configuration in a custom preset", width=90) - - # padding + def _setup_frame_column_weight(self): self.frame.grid_columnconfigure(5, weight=1) - # model type - components.options_kv( - master=self.frame, - row=0, - column=6, - values=[ #TODO simplify - ("SD1.5", ModelType.STABLE_DIFFUSION_15), - ("SD1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), - ("SD2.0", ModelType.STABLE_DIFFUSION_20), - ("SD2.0 Inpainting", ModelType.STABLE_DIFFUSION_20_INPAINTING), - ("SD2.1", ModelType.STABLE_DIFFUSION_21), - ("SD3", ModelType.STABLE_DIFFUSION_3), - ("SD3.5", ModelType.STABLE_DIFFUSION_35), - ("SDXL", ModelType.STABLE_DIFFUSION_XL_10_BASE), - ("SDXL Inpainting", ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING), - ("Wuerstchen v2", ModelType.WUERSTCHEN_2), - ("Stable Cascade", ModelType.STABLE_CASCADE_1), - ("PixArt Alpha", ModelType.PIXART_ALPHA), - ("PixArt Sigma", ModelType.PIXART_SIGMA), - ("Flux Dev.1", ModelType.FLUX_DEV_1), - ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), - ("Flux 2 [Dev, Klein]", ModelType.FLUX_2), - ("Sana", ModelType.SANA), - ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), - ("HiDream Full", ModelType.HI_DREAM_FULL), - ("Chroma1", ModelType.CHROMA_1), - ("QwenImage", ModelType.QWEN), - ("Z-Image", ModelType.Z_IMAGE), - ("Ernie Image", ModelType.ERNIE), - ], - ui_state=self.ui_state, - var_name="model_type", - command=self.__change_model_type, - ) - - def __create_training_method(self): - if self.training_method: - self.training_method.destroy() - - values = [] - #TODO simplify - if self.train_config.model_type.is_stable_diffusion(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ("Embedding", TrainingMethod.EMBEDDING), - ("Fine Tune VAE", TrainingMethod.FINE_TUNE_VAE), - ] - elif self.train_config.model_type.is_stable_diffusion_3() \ - or self.train_config.model_type.is_stable_diffusion_xl() \ - or self.train_config.model_type.is_wuerstchen() \ - or self.train_config.model_type.is_pixart() \ - or self.train_config.model_type.is_flux_1() \ - or self.train_config.model_type.is_sana() \ - or self.train_config.model_type.is_hunyuan_video() \ - or self.train_config.model_type.is_hi_dream() \ - or self.train_config.model_type.is_chroma(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ("Embedding", TrainingMethod.EMBEDDING), - ] - elif self.train_config.model_type.is_qwen() \ - or self.train_config.model_type.is_z_image() \ - or self.train_config.model_type.is_flux_2() \ - or self.train_config.model_type.is_ernie(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ] - - # training method - self.training_method = components.options_kv( - master=self.frame, - row=0, - column=7, - values=values, - ui_state=self.ui_state, - var_name="training_method", - command=self.change_training_method_callback, - ) - - def __change_model_type(self, model_type: ModelType): - self.change_model_type_callback(model_type) - self.__create_training_method() - - def __create_configs_dropdown(self): - if self.configs_dropdown is not None: - self.configs_dropdown.grid_forget() - - self.configs_dropdown = components.options_kv( - self.frame, 0, 1, self.configs, self.config_ui_state, "config_name", self.__load_current_config - ) - - def __load_available_config_names(self): - if os.path.isdir(self.dir): - for path in os.listdir(self.dir): - if path != "#.json": - path = path_util.canonical_join(self.dir, path) - if path.endswith(".json") and os.path.isfile(path): - name = os.path.basename(path) - name = os.path.splitext(name)[0] - self.configs.append((name, path)) - self.configs.sort() - - def __save_to_file(self, name) -> str: - name = path_util.safe_filename(name) - path = path_util.canonical_join("training_presets", f"{name}.json") - - write_json_atomic(path, self.train_config.to_settings_dict(secrets=False)) - - return path - - def __save_secrets(self, path) -> str: - write_json_atomic(path, self.train_config.secrets.to_dict()) - return path - - def open_wiki(self): - webbrowser.open("https://github.com/Nerogar/OneTrainer/wiki", new=0, autoraise=False) - - def __save_new_config(self, name): - path = self.__save_to_file(name) - - is_new_config = name not in [x[0] for x in self.configs] - - if is_new_config: - self.configs.append((name, path)) - self.configs.sort() - - if self.config_ui_data["config_name"] != path_util.canonical_join(self.dir, f"{name}.json"): - self.config_ui_state.get_var("config_name").set(path_util.canonical_join(self.dir, f"{name}.json")) - - if is_new_config: - self.__create_configs_dropdown() - - def __save_config(self): - default_value = self.configs_dropdown.get() - while default_value.startswith('#'): - default_value = default_value[1:] + def _forget_dropdown(self): + self.configs_dropdown.grid_forget() + def _show_save_dialog(self, default_value: str, callback): dialogs.StringInputDialog( parent=self.master, title="name", question="Config Name", - callback=self.__save_new_config, + callback=callback, default_value=default_value, - validate_callback=lambda x: not x.startswith("#") + validate_callback=lambda x: not x.startswith("#"), ) - - def __load_current_config(self, filename): - try: - basename = os.path.basename(filename) - is_built_in_preset = basename.startswith("#") and basename != "#.json" - - with open(filename, "r") as f: - loaded_dict = json.load(f) - default_config = TrainConfig.default_values() - if is_built_in_preset: - # always assume built-in configs are saved in the most recent version - loaded_dict["__version"] = default_config.config_version - loaded_config = default_config.from_dict(loaded_dict).to_unpacked_config() - - with suppress(FileNotFoundError), open("secrets.json", "r") as f: - secrets_dict=json.load(f) - loaded_config.secrets = SecretsConfig.default_values().from_dict(secrets_dict) - - self.train_config.from_dict(loaded_config.to_dict()) - self.ui_state.update(loaded_config) - - optimizer_config = change_optimizer(self.train_config) - self.ui_state.get_var("optimizer").update(optimizer_config) - - self.load_preset_callback() - except FileNotFoundError: - pass - except Exception: - print(traceback.format_exc()) - - def __remove_config(self): - # TODO - pass - - def save_default(self): - self.__save_to_file("#") - self.__save_secrets("secrets.json") diff --git a/modules/ui/CtkTrainUIView.py b/modules/ui/CtkTrainUIView.py index ba90d2e64..d3c1a70b5 100644 --- a/modules/ui/CtkTrainUIView.py +++ b/modules/ui/CtkTrainUIView.py @@ -1,51 +1,40 @@ import ctypes -import datetime -import json -import os import platform -import subprocess -import sys -import threading -import time -import traceback -import webbrowser from collections.abc import Callable from contextlib import suppress from pathlib import Path from tkinter import filedialog, messagebox -import scripts.generate_debug_report -from modules.ui.AdditionalEmbeddingsTab import AdditionalEmbeddingsTab -from modules.ui.CaptionUI import CaptionUI -from modules.ui.CloudTab import CloudTab -from modules.ui.ConceptTab import ConceptTab -from modules.ui.ConvertModelUI import ConvertModelUI -from modules.ui.LoraTab import LoraTab -from modules.ui.ModelTab import ModelTab -from modules.ui.ProfilingWindow import ProfilingWindow -from modules.ui.SampleWindow import SampleWindow -from modules.ui.SamplingTab import SamplingTab -from modules.ui.TopBar import TopBar -from modules.ui.TrainingTab import TrainingTab -from modules.ui.VideoToolUI import VideoToolUI -from modules.util import create -from modules.util.callbacks.TrainCallbacks import TrainCallbacks -from modules.util.commands.TrainCommands import TrainCommands +from modules.ui.AdditionalEmbeddingsTabController import AdditionalEmbeddingsTabController +from modules.ui.BaseTrainUIView import BaseTrainUIView +from modules.ui.CloudTabController import CloudTabController +from modules.ui.ConceptTabController import ConceptTabController +from modules.ui.CtkAdditionalEmbeddingsTabView import CtkAdditionalEmbeddingsTabView +from modules.ui.CtkCaptionUIView import CtkCaptionUIView +from modules.ui.CtkCloudTabView import CtkCloudTabView +from modules.ui.CtkConceptTabView import CtkConceptTabView +from modules.ui.CtkConvertModelUIView import CtkConvertModelUIView +from modules.ui.CtkLoraTabView import CtkLoraTabView +from modules.ui.CtkModelTabView import CtkModelTabView +from modules.ui.CtkProfilingWindowView import CtkProfilingWindowView +from modules.ui.CtkSampleWindowView import CtkSampleWindowView +from modules.ui.CtkSamplingTabView import CtkSamplingTabView +from modules.ui.CtkTopBarView import CtkTopBarView +from modules.ui.CtkTrainingTabView import CtkTrainingTabView +from modules.ui.CtkVideoToolUIView import CtkVideoToolUIView +from modules.ui.LoraTabController import LoraTabController +from modules.ui.ModelTabController import ModelTabController +from modules.ui.ProfilingWindowController import ProfilingWindowController +from modules.ui.SamplingTabController import SamplingTabController +from modules.ui.TopBarController import TopBarController +from modules.ui.TrainingTabController import TrainingTabController +from modules.ui.TrainUIController import TrainUIController from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.DataType import DataType -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.PathIOType import PathIOType from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.torch_util import torch_gc -from modules.util.TrainProgress import TrainProgress -from modules.util.ui import components +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState -from modules.util.ui.validation import flush_and_validate_all - -import torch import customtkinter as ctk from customtkinter import AppearanceModeTracker @@ -57,14 +46,13 @@ # 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 CtkTrainUIView(BaseTrainUIView, ctk.CTk): set_step_progress: Callable[[int, int], None] set_epoch_progress: Callable[[int, int], None] status_label: ctk.CTkLabel | None training_button: ctk.CTkButton | None - training_callbacks: TrainCallbacks | None - training_commands: TrainCommands | None _TRAIN_BUTTON_STYLES = { "idle": { @@ -93,7 +81,8 @@ class TrainUI(ctk.CTk): } def __init__(self): - super().__init__() + ctk.CTk.__init__(self) + BaseTrainUIView.__init__(self, ctk_components) self.title("OneTrainer") self.geometry("1100x740") @@ -105,7 +94,10 @@ def __init__(self): ctk.set_default_color_theme("blue") self.train_config = TrainConfig.default_values() - self.ui_state = UIState(self, self.train_config) + self.ui_state = CtkUIState(self, self.train_config) + + self.controller = TrainUIController(self.train_config) + self.controller.view = self self.grid_rowconfigure(0, weight=0) self.grid_rowconfigure(1, weight=1) @@ -128,80 +120,120 @@ def __init__(self): self.content_frame(self) self.bottom_bar(self) - self.training_thread = None - self.training_callbacks = None - self.training_commands = None + self.controller._check_start_always_on_tensorboard() - self.always_on_tensorboard_subprocess = None - self.current_workspace_dir = self.train_config.workspace_dir - self._check_start_always_on_tensorboard() - - self.workspace_dir_trace_id = self.ui_state.add_var_trace("workspace_dir", self._on_workspace_dir_change_trace) + self.workspace_dir_trace_id = self.ui_state.add_var_trace("workspace_dir", self.controller._on_workspace_dir_change_trace) # Persistent profiling window. - self.profiling_window = ProfilingWindow(self) + self._profiling_controller = ProfilingWindowController() + self.profiling_window = self._profiling_controller.create_window(self, CtkProfilingWindowView) self.protocol("WM_DELETE_WINDOW", self.__close) def __close(self): self.top_bar_component.save_default() - self._stop_always_on_tensorboard() + self.controller._stop_always_on_tensorboard() if hasattr(self, 'workspace_dir_trace_id'): self.ui_state.remove_var_trace("workspace_dir", self.workspace_dir_trace_id) self.quit() + # --- BaseTrainUIView abstract method implementations --- + + def on_update_status(self, status: str): + self.status_label.configure(text=status) + + def on_training_started(self): + self._set_training_button_style("running") + + def on_training_stopped(self, error_caught: bool): + self.eta_label.configure(text="") + self._set_training_button_style("idle") + + def on_training_stopping(self): + self._set_training_button_style("stopping") + + def on_update_progress(self, epoch_step: int, max_step: int, epoch: int, max_epoch: int, eta_str: str | None): + self.set_step_progress(epoch_step, max_step) + self.set_epoch_progress(epoch, max_epoch) + if eta_str is not None: + self.eta_label.configure(text=f"ETA: {eta_str}") + else: + self.eta_label.configure(text="") + + def schedule_on_main_thread(self, fn: Callable): + self.after(0, fn) + + def get_cloud_reattach(self) -> bool: + return self.cloud_tab.reattach + + 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() + + def show_validation_errors(self, errors: list[str]): + bullet_list = "\n".join(f"• {e}" for e in errors) + messagebox.showerror( + "Cannot Start Training", + f"Please fix the following errors before training:\n\n{bullet_list}", + ) + + def open_dataset_tool(self): + self.wait_window(self.controller.open_dataset_tool(self, CtkCaptionUIView)) + + def open_video_tool(self): + self.wait_window(self.controller.open_video_tool(self, CtkVideoToolUIView)) + + def open_convert_model_tool(self): + self.wait_window(self.controller.open_convert_model_tool(self, CtkConvertModelUIView)) + + def open_sampling_tool(self): + self.controller.open_sampling_tool(self, CtkSampleWindowView) + + def open_manual_sample_window(self): + self.controller.open_manual_sample_window(self, CtkSampleWindowView) + + def wait_window(self, window): + ctk.CTk.wait_window(self, window) + + def show_window(self, window): + window.focus_set() + + def connect_window_closed(self, window, callback): + window.bind("", lambda _: callback()) + + # --- CTK layout and frame builders --- + + def _set_icon(self): + """Set the window icon safely after window is ready""" + set_window_icon(self) + def top_bar(self, master): - return TopBar( + return CtkTopBarView( master, - self.train_config, + TopBarController(self.train_config), self.ui_state, self.change_model_type, self.change_training_method, self.load_preset, ) - def _set_icon(self): - """Set the window icon safely after window is ready""" - set_window_icon(self) - def bottom_bar(self, master): frame = ctk.CTkFrame(master=master, corner_radius=0) frame.grid(row=2, column=0, sticky="nsew") - self.set_step_progress, self.set_epoch_progress = components.double_progress(frame, 0, 0, "step", "epoch") - # status + ETA container - self.status_frame = ctk.CTkFrame(frame, corner_radius=0, fg_color="transparent") - self.status_frame.grid(row=0, column=1, sticky="w") - self.status_frame.grid_rowconfigure(0, weight=0) - self.status_frame.grid_rowconfigure(1, weight=0) - self.status_frame.grid_columnconfigure(0, weight=1) - - self.status_label = components.label(self.status_frame, 0, 0, "", pad=0, - tooltip="Current status of the training run") - self.eta_label = components.label(self.status_frame, 1, 0, "", pad=0) + status_frame = ctk.CTkFrame(frame, corner_radius=0, fg_color="transparent") + status_frame.grid(row=0, column=1, sticky="w") + status_frame.grid_rowconfigure(0, weight=0) + status_frame.grid_rowconfigure(1, weight=0) + status_frame.grid_columnconfigure(0, weight=1) # padding frame.grid_columnconfigure(2, weight=1) - - # export button - self.export_button = components.button(frame, 0, 3, "Export", self.export_training, - width=60, padx=5, pady=(15, 0), - tooltip="Export the current configuration as a script to run without a UI") - - # debug button - components.button(frame, 0, 4, "Debug", self.generate_debug_package, - width=60, padx=(5, 25), pady=(15, 0), - tooltip="Generate a zip file with config.json, debug_report.log and settings diff, use this to report bugs or issues") - - # tensorboard button - components.button(frame, 0, 5, "Tensorboard", self.open_tensorboard, - width=100, padx=(0, 5), pady=(15, 0)) - - # training button - self.training_button = components.button(frame, 0, 6, "Start Training", self.start_training, - padx=(5, 20), pady=(15, 0)) + self.build_bottom_bar_content(frame, status_frame, self.controller, self.ui_state) self._set_training_button_style("idle") # centralized styling return frame @@ -237,113 +269,12 @@ def create_general_tab(self, master): frame.grid_columnconfigure(1, weight=1) frame.grid_columnconfigure(2, weight=0) frame.grid_columnconfigure(3, weight=1) - - # workspace dir - components.label(frame, 0, 0, "Workspace Directory", - tooltip="The directory where all files of this training run are saved") - components.path_entry(frame, 0, 1, self.ui_state, "workspace_dir", mode="dir", command=self._on_workspace_dir_change) - - # cache dir - components.label(frame, 0, 2, "Cache Directory", - tooltip="The directory where cached data is saved") - components.path_entry(frame, 0, 3, self.ui_state, "cache_dir", mode="dir") - - # continue from previous backup - components.label(frame, 2, 0, "Continue from last backup", - tooltip="Automatically continues training from the last backup saved in /backup") - components.switch(frame, 2, 1, self.ui_state, "continue_last_backup") - - # only cache - components.label(frame, 2, 2, "Only Cache", - tooltip="Only populate the cache, without any training") - components.switch(frame, 2, 3, self.ui_state, "only_cache") - - # TODO: In Phase 4 rework the general tab. - # prevent overwrites - components.label(frame, 3, 0, "Prevent Overwrites", - tooltip="When enabled, output paths that already exist on disk will be flagged as invalid to avoid accidental overwrites") - components.switch(frame, 3, 1, self.ui_state, "prevent_overwrites") - - # debug - components.label(frame, 4, 0, "Debug mode", - tooltip="Save debug information during the training into the debug directory") - components.switch(frame, 4, 1, self.ui_state, "debug_mode") - - components.label(frame, 4, 2, "Debug Directory", - tooltip="The directory where debug data is saved") - components.path_entry(frame, 4, 3, self.ui_state, "debug_dir", mode="dir", io_type=PathIOType.OUTPUT) - - # 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", - tooltip="Port to use for Tensorboard link") - components.entry(frame, 7, 3, self.ui_state, "tensorboard_port") - - - # validation - components.label(frame, 8, 0, "Validation", - tooltip="Enable validation steps and add new graph in tensorboard") - components.switch(frame, 8, 1, self.ui_state, "validation") - - components.label(frame, 8, 2, "Validate after", - tooltip="The interval used when validate training") - components.time_entry(frame, 8, 3, self.ui_state, "validate_after", "validate_after_unit") - - # device - components.label(frame, 10, 0, "Dataloader Threads", - tooltip="Number of threads used for the data loader. Increase if your GPU has room during caching, decrease if it's going out of memory during caching.") - components.entry(frame, 10, 1, self.ui_state, "dataloader_threads", required=True) - - components.label(frame, 11, 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", required=True) - - components.label(frame, 12, 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", - 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.label(frame, 13, 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, - "gradient_reduce_precision") - - components.label(frame, 13, 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.label(frame, 14, 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)", - 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.label(frame, 15, 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") - + self.build_general_tab_content(frame, self.controller, self.ui_state) frame.pack(fill="both", expand=1) return frame def create_model_tab(self, master): - return ModelTab(master, self.train_config, self.ui_state) + return CtkModelTabView(master, ModelTabController(self.train_config), self.ui_state) def create_data_tab(self, master): frame = ctk.CTkScrollableFrame(master, fg_color="transparent") @@ -352,77 +283,35 @@ def create_data_tab(self, master): frame.grid_columnconfigure(2, minsize=50) frame.grid_columnconfigure(3, weight=0) frame.grid_columnconfigure(4, weight=1) - - # aspect ratio bucketing - components.label(frame, 0, 0, "Aspect Ratio Bucketing", - tooltip="Aspect ratio bucketing enables training on images with different aspect ratios") - components.switch(frame, 0, 1, self.ui_state, "aspect_ratio_bucketing") - - # latent caching - components.label(frame, 1, 0, "Latent Caching", - tooltip="Caching of intermediate training data that can be re-used between epochs") - components.switch(frame, 1, 1, self.ui_state, "latent_caching") - - # clear cache before training - components.label(frame, 2, 0, "Clear cache before training", - tooltip="Clears the cache directory before starting to train. Only disable this if you want to continue using the same cached data. Disabling this can lead to errors, if other settings are changed during a restart") - components.switch(frame, 2, 1, self.ui_state, "clear_cache_before_training") - + self.build_data_tab_content(frame, self.controller, self.ui_state) frame.pack(fill="both", expand=1) return frame def create_concepts_tab(self, master): - return ConceptTab(master, self.train_config, self.ui_state) + return CtkConceptTabView(master, ConceptTabController(self.train_config), self.ui_state) - def create_training_tab(self, master) -> TrainingTab: - return TrainingTab(master, self.train_config, self.ui_state) + def create_training_tab(self, master) -> CtkTrainingTabView: + return CtkTrainingTabView(master, TrainingTabController(self.train_config), self.ui_state) - def create_cloud_tab(self, master) -> CloudTab: - return CloudTab(master, self.train_config, self.ui_state,parent=self) + def create_cloud_tab(self, master) -> CtkCloudTabView: + return CtkCloudTabView(master, CloudTabController(self.train_config, parent=self), self.ui_state) def create_sampling_tab(self, master): master.grid_rowconfigure(0, weight=0) master.grid_rowconfigure(1, weight=1) master.grid_columnconfigure(0, weight=1) - # sample after top_frame = ctk.CTkFrame(master=master, corner_radius=0) top_frame.grid(row=0, column=0, sticky="nsew") sub_frame = ctk.CTkFrame(master=top_frame, corner_radius=0, fg_color="transparent") sub_frame.grid(row=1, column=0, sticky="nsew", columnspan=6) - components.label(top_frame, 0, 0, "Sample After", - tooltip="The interval used when automatically sampling from the model during training") - components.time_entry(top_frame, 0, 1, self.ui_state, "sample_after", "sample_after_unit") - - components.label(top_frame, 0, 2, "Skip First", - tooltip="Start sampling automatically after this interval has elapsed.") - components.entry(top_frame, 0, 3, self.ui_state, "sample_skip_first", width=50, sticky="nw") - - components.label(top_frame, 0, 4, "Format", - tooltip="File Format used when saving samples") - components.options_kv(top_frame, 0, 5, [ - ("PNG", ImageFormat.PNG), - ("JPG", ImageFormat.JPG), - ], self.ui_state, "sample_image_format") - - components.button(top_frame, 0, 6, "sample now", self.sample_now) - - components.button(top_frame, 0, 7, "manual sample", self.open_manual_sample_window ) - - components.label(sub_frame, 0, 0, "Non-EMA Sampling", - tooltip="Whether to include non-ema sampling when using ema.") - components.switch(sub_frame, 0, 1, self.ui_state, "non_ema_sampling") - - components.label(sub_frame, 0, 2, "Samples to Tensorboard", - tooltip="Whether to include sample images in the Tensorboard output.") - components.switch(sub_frame, 0, 3, self.ui_state, "samples_to_tensorboard") + self.build_sampling_tab_header(top_frame, sub_frame, self.controller, self.ui_state) - # table frame = ctk.CTkFrame(master=master, corner_radius=0) frame.grid(row=1, column=0, sticky="nsew") - return SamplingTab(frame, self.train_config, self.ui_state) + return CtkSamplingTabView(frame, SamplingTabController(self.train_config), self.ui_state) def create_backup_tab(self, master): frame = ctk.CTkScrollableFrame(master, fg_color="transparent") @@ -431,48 +320,7 @@ def create_backup_tab(self, master): frame.grid_columnconfigure(2, minsize=50) frame.grid_columnconfigure(3, weight=0) frame.grid_columnconfigure(4, weight=1) - - # backup after - components.label(frame, 0, 0, "Backup After", - tooltip="The interval used when automatically creating model backups during training") - components.time_entry(frame, 0, 1, self.ui_state, "backup_after", "backup_after_unit") - - # backup now - components.button(frame, 0, 3, "backup now", self.backup_now) - - # rolling backup - components.label(frame, 1, 0, "Rolling Backup", - tooltip="If rolling backups are enabled, older backups are deleted automatically") - components.switch(frame, 1, 1, self.ui_state, "rolling_backup") - - # rolling backup count - components.label(frame, 1, 3, "Rolling Backup Count", - tooltip="Defines the number of backups to keep if rolling backups are enabled") - components.entry(frame, 1, 4, self.ui_state, "rolling_backup_count") - - # backup before save - components.label(frame, 2, 0, "Backup Before Save", - tooltip="Create a full backup before saving the final model") - components.switch(frame, 2, 1, self.ui_state, "backup_before_save") - - # save after - components.label(frame, 3, 0, "Save Every", - tooltip="The interval used when automatically saving the model during training") - components.time_entry(frame, 3, 1, self.ui_state, "save_every", "save_every_unit") - - # save now - components.button(frame, 3, 3, "save now", self.save_now) - - # skip save - components.label(frame, 4, 0, "Skip First", - tooltip="Start saving automatically after this interval has elapsed") - components.entry(frame, 4, 1, self.ui_state, "save_skip_first", width=50, sticky="nw") - - # save filename prefix - components.label(frame, 5, 0, "Save Filename Prefix", - tooltip="The prefix for filenames used when saving the model during training") - components.entry(frame, 5, 1, self.ui_state, "save_filename_prefix") - + self.build_backup_tab_content(frame, self.controller, self.ui_state) frame.pack(fill="both", expand=1) return frame @@ -483,48 +331,12 @@ def embedding_tab(self, master): frame.grid_columnconfigure(2, minsize=50) frame.grid_columnconfigure(3, weight=0) frame.grid_columnconfigure(4, weight=1) - - # 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.path_entry( - frame, 0, 1, self.ui_state, "embedding.model_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # token count - components.label(frame, 1, 0, "Token count", - tooltip="The token count used when creating a new embedding. Leave empty to auto detect from the initial embedding text.") - components.entry(frame, 1, 1, self.ui_state, "embedding.token_count") - - # initial embedding text - components.label(frame, 2, 0, "Initial embedding text", - tooltip="The initial embedding text used when creating a new embedding") - components.entry(frame, 2, 1, self.ui_state, "embedding.initial_embedding_text") - - # embedding weight dtype - components.label(frame, 3, 0, "Embedding Weight Data Type", - tooltip="The Embedding weight data type used for training. This can reduce memory consumption, but reduces precision") - components.options_kv(frame, 3, 1, [ - ("float32", DataType.FLOAT_32), - ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "embedding_weight_dtype") - - # placeholder - components.label(frame, 4, 0, "Placeholder", - tooltip="The placeholder used when using the embedding in a prompt") - components.entry(frame, 4, 1, self.ui_state, "embedding.placeholder") - - # output embedding - components.label(frame, 5, 0, "Output embedding", - tooltip="Output embeddings are calculated at the output of the text encoder, not the input. This can improve results for larger text encoders and lower VRAM usage.") - components.switch(frame, 5, 1, self.ui_state, "embedding.is_output_embedding") - + self.build_embedding_tab_content(frame, self.controller, self.ui_state) frame.pack(fill="both", expand=1) return frame def create_additional_embeddings_tab(self, master): - return AdditionalEmbeddingsTab(master, self.train_config, self.ui_state) + return CtkAdditionalEmbeddingsTabView(master, AdditionalEmbeddingsTabController(self.train_config), self.ui_state) def create_tools_tab(self, master): frame = ctk.CTkScrollableFrame(master, fg_color="transparent") @@ -533,34 +345,13 @@ def create_tools_tab(self, master): frame.grid_columnconfigure(2, minsize=50) frame.grid_columnconfigure(3, weight=0) frame.grid_columnconfigure(4, weight=1) - - # dataset - components.label(frame, 0, 0, "Dataset Tools", - tooltip="Open the captioning tool") - components.button(frame, 0, 1, "Open", self.open_dataset_tool) - - # video tools - components.label(frame, 1, 0, "Video Tools", - tooltip="Open the video tools") - components.button(frame, 1, 1, "Open", self.open_video_tool) - - # convert model - components.label(frame, 2, 0, "Convert Model Tools", - tooltip="Open the model conversion tool") - components.button(frame, 2, 1, "Open", self.open_convert_model_tool) - - # sample - components.label(frame, 3, 0, "Sampling Tool", - tooltip="Open the model sampling tool") - components.button(frame, 3, 1, "Open", self.open_sampling_tool) - - components.label(frame, 4, 0, "Profiling Tool", - tooltip="Open the profiling tools.") - components.button(frame, 4, 1, "Open", self.open_profiling_tool) - + self.build_tools_tab_content(frame, self.controller, self.ui_state) frame.pack(fill="both", expand=1) return frame + def open_profiling_tool(self): + self.profiling_window.deiconify() + def change_model_type(self, model_type: ModelType): if self.model_tab: self.model_tab.refresh_ui() @@ -585,7 +376,7 @@ def change_training_method(self, training_method: TrainingMethod): self.tabview.delete("embedding") if training_method == TrainingMethod.LORA and "LoRA" not in self.tabview._tab_dict: - self.lora_tab = LoraTab(self.tabview.add("LoRA"), self.train_config, self.ui_state) + self.lora_tab = CtkLoraTabView(self.tabview.add("LoRA"), LoraTabController(self.train_config), self.ui_state) if training_method == TrainingMethod.EMBEDDING and "embedding" not in self.tabview._tab_dict: self.embedding_tab(self.tabview.add("embedding")) @@ -596,294 +387,27 @@ def load_preset(self): if self.additional_embeddings_tab: self.additional_embeddings_tab.refresh_ui() - def open_tensorboard(self): - webbrowser.open("http://localhost:" + str(self.train_config.tensorboard_port), new=0, autoraise=False) - - def _calculate_eta_string(self, train_progress: TrainProgress, max_step: int, max_epoch: int) -> str | None: - spent_total = time.monotonic() - self.start_time - steps_done = train_progress.epoch * max_step + train_progress.epoch_step - remaining_steps = (max_epoch - train_progress.epoch - 1) * max_step + (max_step - train_progress.epoch_step) - total_eta = spent_total / steps_done * remaining_steps - - if train_progress.global_step <= 30: - return "Estimating ..." - - td = datetime.timedelta(seconds=total_eta) - days = td.days - hours, remainder = divmod(td.seconds, 3600) - minutes, seconds = divmod(remainder, 60) - if days > 0: - return f"{days}d {hours}h" - elif hours > 0: - return f"{hours}h {minutes}m" - elif minutes > 0: - return f"{minutes}m {seconds}s" - else: - return f"{seconds}s" - - def set_eta_label(self, train_progress: TrainProgress, max_step: int, max_epoch: int): - eta_str = self._calculate_eta_string(train_progress, max_step, max_epoch) - if eta_str is not None: - self.eta_label.configure(text=f"ETA: {eta_str}") - else: - self.eta_label.configure(text="") - - def delete_eta_label(self): - self.eta_label.configure(text="") - - def on_update_train_progress(self, train_progress: TrainProgress, max_step: int, max_epoch: int): - self.set_step_progress(train_progress.epoch_step, max_step) - self.set_epoch_progress(train_progress.epoch, max_epoch) - self.set_eta_label(train_progress, max_step, max_epoch) - - def on_update_status(self, status: str): - self.status_label.configure(text=status) - - def open_dataset_tool(self): - window = CaptionUI(self, None, False) - self.wait_window(window) - - def open_video_tool(self): - window = VideoToolUI(self) - self.wait_window(window) - - def open_convert_model_tool(self): - window = ConvertModelUI(self) - self.wait_window(window) - - def open_sampling_tool(self): - if not self.training_callbacks and not self.training_commands: - window = SampleWindow( - self, - use_external_model=False, - train_config=self.train_config, - ) - self.wait_window(window) - torch_gc() - - def open_profiling_tool(self): - self.profiling_window.deiconify() - - def generate_debug_package(self): - zip_path = filedialog.askdirectory( - initialdir=".", - title="Select Directory to Save Debug Package" - ) - - if not zip_path: + def _set_training_button_style(self, mode: str): + if not self.training_button: return - - zip_path = Path(zip_path) / "OneTrainer_debug_report.zip" - - self.on_update_status("Generating debug package...") - - try: - config_json_string = json.dumps(self.train_config.to_pack_dict(secrets=False)) - scripts.generate_debug_report.create_debug_package(str(zip_path), config_json_string) - self.on_update_status(f"Debug package saved to {zip_path.name}") - except Exception as e: - traceback.print_exc() - self.on_update_status(f"Error generating debug package: {e}") - - - def open_manual_sample_window (self): - training_callbacks = self.training_callbacks - training_commands = self.training_commands - - if training_callbacks and training_commands: - window = SampleWindow( - self, - train_config=self.train_config, - use_external_model=True, - callbacks=training_callbacks, - commands=training_commands, - ) - self.wait_window(window) - training_callbacks.set_on_sample_custom() - - def __training_thread_function(self): - error_caught = False - - self.training_callbacks = TrainCallbacks( - on_update_train_progress=self.on_update_train_progress, - on_update_status=self.on_update_status, - ) - - trainer = create.create_trainer(self.train_config, self.training_callbacks, self.training_commands, reattach=self.cloud_tab.reattach) - try: - trainer.start() - if self.train_config.cloud.enabled: - self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) - - self.start_time = time.monotonic() - trainer.train() - except Exception: - if self.train_config.cloud.enabled: - self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) - error_caught = True - traceback.print_exc() - - trainer.end() - - # clear gpu memory - del trainer - - 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() - - # 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) - - def start_training(self): - if self.training_thread is None: - self.save_default() - - # --- pre-training validation gate --- - errors = flush_and_validate_all() - - if errors: - bullet_list = "\n".join(f"• {e}" for e in errors) - messagebox.showerror( - "Cannot Start Training", - f"Please fix the following errors before training:\n\n{bullet_list}", - ) - return - - 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() - - self.training_thread = threading.Thread(target=self.__training_thread_function) - self.training_thread.start() - else: - self._set_training_button_stopping() - self.on_update_status("Stopping ...") - self.training_commands.stop() - - 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() + style = self._TRAIN_BUTTON_STYLES.get(mode) + if not style: + return + self.training_button.configure(**style) def export_training(self): file_path = filedialog.asksaveasfilename(filetypes=[ ("All Files", "*.*"), ("json", "*.json"), ], initialdir=".", initialfile="config.json") - if file_path: - with open(file_path, "w") as f: - json.dump(self.train_config.to_pack_dict(secrets=False), f, indent=4) - - def sample_now(self): - train_commands = self.training_commands - if train_commands: - train_commands.sample_default() - - def backup_now(self): - train_commands = self.training_commands - if train_commands: - train_commands.backup() - - def save_now(self): - train_commands = self.training_commands - 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") - - 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", - ] - - if self.train_config.tensorboard_expose: - tensorboard_args.append("--bind_all") - - 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() + self.controller.export_training(file_path) - def _set_training_button_style(self, mode: str): - if not self.training_button: - return - style = self._TRAIN_BUTTON_STYLES.get(mode) - if not style: + def generate_debug_package(self): + zip_path = filedialog.askdirectory( + initialdir=".", + title="Select Directory to Save Debug Package" + ) + if not zip_path: return - self.training_button.configure(**style) - - def _set_training_button_idle(self): - self._set_training_button_style("idle") - - def _set_training_button_running(self): - self._set_training_button_style("running") - - def _set_training_button_stopping(self): - self._set_training_button_style("stopping") + self.controller.generate_debug_package(Path(zip_path) / "OneTrainer_debug_report.zip") diff --git a/modules/ui/CtkTrainingTabView.py b/modules/ui/CtkTrainingTabView.py index bcca11ae9..bc29488dd 100644 --- a/modules/ui/CtkTrainingTabView.py +++ b/modules/ui/CtkTrainingTabView.py @@ -1,40 +1,27 @@ -from modules.ui.OffloadingWindow import OffloadingWindow -from modules.ui.OptimizerParamsWindow import OptimizerParamsWindow -from modules.ui.SchedulerParamsWindow import SchedulerParamsWindow -from modules.ui.TimestepDistributionWindow import TimestepDistributionWindow -from modules.util import create -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.DataType import DataType -from modules.util.enum.EMAMode import EMAMode -from modules.util.enum.GradientCheckpointingMethod import GradientCheckpointingMethod -from modules.util.enum.LearningRateScaler import LearningRateScaler -from modules.util.enum.LearningRateScheduler import LearningRateScheduler -from modules.util.enum.LossScaler import LossScaler -from modules.util.enum.LossWeight import LossWeight -from modules.util.enum.Optimizer import Optimizer -from modules.util.enum.TimestepDistribution import TimestepDistribution -from modules.util.optimizer_util import change_optimizer -from modules.util.ui import components -from modules.util.ui.UIState import UIState -from modules.util.ui.validation_helpers import check_range, validate_resolution -import customtkinter as ctk +from modules.ui.BaseTrainingTabView import BaseTrainingTabView +from modules.ui.CtkOffloadingWindowView import CtkOffloadingWindowView +from modules.ui.CtkOptimizerParamsWindowView import CtkOptimizerParamsWindowView +from modules.ui.CtkSchedulerParamsWindowView import CtkSchedulerParamsWindowView +from modules.ui.CtkTimestepDistributionWindowView import CtkTimestepDistributionWindowView +from modules.ui.TrainingTabController import TrainingTabController +from modules.util.ui import ctk_components +import customtkinter as ctk -class TrainingTab: - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__() +class CtkTrainingTabView(BaseTrainingTabView): + def __init__(self, master, controller: TrainingTabController, ui_state): + BaseTrainingTabView.__init__(self, ctk_components) self.master = master - self.train_config = train_config + self.controller = controller self.ui_state = ui_state + self.scroll_frame = None master.grid_rowconfigure(0, weight=1) master.grid_columnconfigure(0, weight=1) - self.scroll_frame = None - self.refresh_ui() def refresh_ui(self): @@ -60,797 +47,31 @@ def refresh_ui(self): column_2.grid(row=0, column=2, sticky="nsew") column_2.grid_columnconfigure(0, weight=1) - if self.train_config.model_type.is_stable_diffusion(): - self.__setup_stable_diffusion_ui(column_0, column_1, column_2) - if self.train_config.model_type.is_stable_diffusion_3(): - self.__setup_stable_diffusion_3_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_stable_diffusion_xl(): - self.__setup_stable_diffusion_xl_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_wuerstchen(): - self.__setup_wuerstchen_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_pixart(): - self.__setup_pixart_alpha_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_flux_1(): - self.__setup_flux_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_flux_2(): - self.__setup_flux_2_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_chroma(): - self.__setup_chroma_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_qwen(): - self.__setup_qwen_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_sana(): - self.__setup_sana_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_hunyuan_video(): - self.__setup_hunyuan_video_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_hi_dream(): - self.__setup_hi_dream_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_z_image(): - self.__setup_z_image_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_ernie(): - self.__setup_ernie_ui(column_0, column_1, column_2) - - - def __setup_stable_diffusion_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1) - self.__create_embedding_frame(column_0, 2) - - self.__create_base2_frame(column_1, 0, supports_circular_padding=True) - self.__create_unet_frame(column_1, 1) - self.__create_noise_frame(column_1, 2, supports_generalized_offset_noise=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_stable_diffusion_3_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 3, i=3, supports_include=True) - self.__create_embedding_frame(column_0, 4) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_stable_diffusion_xl_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_n_frame(column_0, 1, i=1) - self.__create_text_encoder_n_frame(column_0, 2, i=2) - self.__create_embedding_frame(column_0, 3) - - self.__create_base2_frame(column_1, 0, supports_circular_padding=True) - self.__create_unet_frame(column_1, 1) - self.__create_noise_frame(column_1, 2, supports_generalized_offset_noise=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_wuerstchen_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1) - self.__create_embedding_frame(column_0, 2) - - self.__create_base2_frame(column_1, 0, supports_circular_padding=True) - self.__create_prior_frame(column_1, 1) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 0) - self.__create_loss_frame(column_2, 1) - self.__create_layer_frame(column_2, 2) - - def __setup_pixart_alpha_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1) - self.__create_embedding_frame(column_0, 2) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2, supports_vb_loss=True) - self.__create_layer_frame(column_2, 3) - - def __setup_flux_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True, supports_sequence_length=True) - self.__create_embedding_frame(column_0, 4) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=True) - self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_flux_2_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False, supports_sequence_length=True) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=True, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_chroma_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1) - self.__create_embedding_frame(column_0, 4) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_qwen_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_z_image_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_ernie_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_sana_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1) - self.__create_embedding_frame(column_0, 2) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_hunyuan_video_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True) - self.__create_embedding_frame(column_0, 4) - - self.__create_base2_frame(column_1, 0, video_training_enabled=True) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=True) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_hi_dream_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 3, i=3, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 4, i=4, supports_include=True, supports_layer_skip=False) - self.__create_embedding_frame(column_0, 5) - - self.__create_base2_frame(column_1, 0, video_training_enabled=True) - self.__create_transformer_frame(column_1, 1) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __create_base_frame(self, master, row): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - - # optimizer - components.label(frame, 0, 0, "Optimizer", - tooltip="The type of optimizer") - components.options_adv(frame, 0, 1, [str(x) for x in list(Optimizer)], self.ui_state, "optimizer.optimizer", - command=self.__restore_optimizer_config, adv_command=self.__open_optimizer_params_window) - - # learning rate scheduler - # Wackiness will ensue when reloading configs if we don't check and clear this first. - if hasattr(self, "lr_scheduler_comp"): - delattr(self, "lr_scheduler_comp") - delattr(self, "lr_scheduler_adv_comp") - components.label(frame, 1, 0, "Learning Rate Scheduler", - tooltip="Learning rate scheduler that automatically changes the learning rate during training") - _, d = components.options_adv(frame, 1, 1, [str(x) for x in list(LearningRateScheduler)], self.ui_state, - "learning_rate_scheduler", command=self.__restore_scheduler_config, - adv_command=self.__open_scheduler_params_window) - self.lr_scheduler_comp = d['component'] - self.lr_scheduler_adv_comp = d['button_component'] - # Initial call requires the presence of self.lr_scheduler_adv_comp. - self.__restore_scheduler_config(self.ui_state.get_var("learning_rate_scheduler").get()) - - # learning rate - components.label(frame, 2, 0, "Learning Rate", - tooltip="The base learning rate") - components.entry(frame, 2, 1, self.ui_state, "learning_rate", required=True) - - # learning rate warmup steps - components.label(frame, 3, 0, "Learning Rate Warmup Steps", - tooltip="The number of steps it takes to gradually increase the learning rate from 0 to the specified learning rate. Values >1 are interpeted as a fixed number of steps, values <=1 are intepreted as a percentage of the total training steps (ex. 0.2 = 20% of the total step count)") - components.entry(frame, 3, 1, self.ui_state, "learning_rate_warmup_steps") - - # learning rate min factor - components.label(frame, 4, 0, "Learning Rate Min Factor", - tooltip="Unit = float. Method = percentage. For a factor of 0.1, the final LR will be 10% of the initial LR. If the initial LR is 1e-4, the final LR will be 1e-5.") - components.entry(frame, 4, 1, self.ui_state, "learning_rate_min_factor", - extra_validate=check_range(lower=0, upper=0.99, message="Learning rate min factor must be between 0 and 0.99")) - - # learning rate cycles - components.label(frame, 5, 0, "Learning Rate Cycles", - tooltip="The number of learning rate cycles. This is only applicable if the learning rate scheduler supports cycles") - components.entry(frame, 5, 1, self.ui_state, "learning_rate_cycles") - - # epochs - components.label(frame, 6, 0, "Epochs", - tooltip="The number of epochs for a full training run") - components.entry(frame, 6, 1, self.ui_state, "epochs", required=True) - - # batch size - components.label(frame, 7, 0, "Local Batch Size", - tooltip="The batch size of one training step. If you use multiple GPUs, this is the batch size of each GPU (local batch size).") - components.entry(frame, 7, 1, self.ui_state, "batch_size", required=True) - - # accumulation steps - components.label(frame, 8, 0, "Accumulation Steps", - tooltip="Number of accumulation steps. Increase this number to trade batch size for training speed") - components.entry(frame, 8, 1, self.ui_state, "gradient_accumulation_steps", required=True) - - # Learning Rate Scaler - components.label(frame, 9, 0, "Learning Rate Scaler", - tooltip="Selects the type of learning rate scaling to use during training. Functionally equated as: LR * SQRT(selection)") - components.options(frame, 9, 1, [str(x) for x in list(LearningRateScaler)], self.ui_state, - "learning_rate_scaler") - - # clip grad norm - components.label(frame, 10, 0, "Clip Grad Norm", - tooltip="Clips the gradient norm. Leave empty to disable gradient clipping.") - components.entry(frame, 10, 1, self.ui_state, "clip_grad_norm") - - def __create_base2_frame(self, master, row, video_training_enabled: bool=False, supports_circular_padding: bool=False): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - row = 0 - - # ema - components.label(frame, row, 0, "EMA", - tooltip="EMA averages the training progress over many steps, better preserving different concepts in big datasets") - components.options(frame, row, 1, [str(x) for x in list(EMAMode)], self.ui_state, "ema") - row += 1 - - # ema decay - components.label(frame, row, 0, "EMA Decay", - tooltip="Decay parameter of the EMA model. Higher numbers will average more steps. For datasets of hundreds or thousands of images, set this to 0.9999. For smaller datasets, set it to 0.999 or even 0.998") - components.entry(frame, row, 1, self.ui_state, "ema_decay", - extra_validate=check_range(lower=0.5, upper=1, - message="EMA decay must be between 0.5 and 1")) - row += 1 - - # ema update step interval - components.label(frame, row, 0, "EMA Update Step Interval", - tooltip="Number of steps between EMA update steps") - components.entry(frame, row, 1, self.ui_state, "ema_update_step_interval") - row += 1 - - # gradient checkpointing - components.label(frame, row, 0, "Gradient checkpointing", - tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") - components.options_adv(frame, row, 1, [str(x) for x in list(GradientCheckpointingMethod)], self.ui_state, - "gradient_checkpointing", adv_command=self.__open_offloading_window) - row += 1 - - # gradient checkpointing layer offloading - components.label(frame, row, 0, "Layer offload fraction", - tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") - components.entry(frame, row, 1, self.ui_state, "layer_offload_fraction") - row += 1 - - # train dtype - components.label(frame, row, 0, "Train Data Type", - tooltip="The mixed precision data type used for training. This can increase training speed, but reduces precision") - components.options_kv(frame, row, 1, [ - ("float32", DataType.FLOAT_32), - ("float16", DataType.FLOAT_16), - ("bfloat16", DataType.BFLOAT_16), - ("tfloat32", DataType.TFLOAT_32), - ], self.ui_state, "train_dtype") - row += 1 - - # fallback train dtype - components.label(frame, row, 0, "Fallback Train Data Type", - tooltip="The mixed precision data type used for training stages that don't support float16 data types. This can increase training speed, but reduces precision") - components.options_kv(frame, row, 1, [ - ("float32", DataType.FLOAT_32), - ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "fallback_train_dtype") - row += 1 - - # autocast cache - components.label(frame, row, 0, "Autocast Cache", - tooltip="Enables the autocast cache. Disabling this reduces memory usage, but increases training time") - components.switch(frame, row, 1, self.ui_state, "enable_autocast_cache") - row += 1 - - # resolution - components.label(frame, row, 0, "Resolution", - tooltip="The resolution used for training. Optionally specify multiple resolutions separated by a comma, or a single exact resolution in the format x") - components.entry(frame, row, 1, self.ui_state, "resolution", required=True, - extra_validate=validate_resolution()) - row += 1 - - # frames - if video_training_enabled: - components.label(frame, row, 0, "Frames", - tooltip="The number of frames used for training.") - components.entry(frame, row, 1, self.ui_state, "frames", required=True) - row += 1 - - # force circular padding - if supports_circular_padding: - components.label(frame, row, 0, "Force Circular Padding", - tooltip="Enables circular padding for all conv layers to better train seamless images") - components.switch(frame, row, 1, self.ui_state, "force_circular_padding") - - def __create_text_encoder_frame(self, master, row, supports_clip_skip=True, supports_training=True, supports_sequence_length=False): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - - if supports_training: - components.label(frame, 0, 0, "Train Text Encoder", - tooltip="Enables training the text encoder model") - components.switch(frame, 0, 1, self.ui_state, "text_encoder.train") + callbacks = { + 'restore_optimizer': lambda *args: self.controller.restore_optimizer_config(self.ui_state), + 'open_optimizer_params': self._open_optimizer_params_window, + 'restore_scheduler': self._restore_scheduler_config, + 'open_scheduler_params': self._open_scheduler_params_window, + 'open_offloading': self._open_offloading_window, + 'open_timestep_distribution': self._open_timestep_distribution_window, + } - # dropout - components.label(frame, 1, 0, "Caption Dropout Probability", - tooltip="The Probability for dropping the text encoder conditioning") - components.entry(frame, 1, 1, self.ui_state, "text_encoder.dropout_probability") + self.build(column_0, column_1, column_2, self.controller, self.ui_state, callbacks) - if supports_training: - # train text encoder epochs - components.label(frame, 2, 0, "Stop Training After", - tooltip="When to stop training the text encoder") - components.time_entry(frame, 2, 1, self.ui_state, "text_encoder.stop_training_after", - "text_encoder.stop_training_after_unit", supports_time_units=False) - - # text encoder learning rate - components.label(frame, 3, 0, "Text Encoder Learning Rate", - tooltip="The learning rate of the text encoder. Overrides the base learning rate") - components.entry(frame, 3, 1, self.ui_state, "text_encoder.learning_rate") - - if supports_clip_skip: - # text encoder layer skip (clip skip) - components.label(frame, 4, 0, "Clip Skip", - tooltip="The number of additional clip layers to skip. 0 = the model default") - components.entry(frame, 4, 1, self.ui_state, "text_encoder_layer_skip") - - if supports_sequence_length: - # text encoder sequence length - components.label(frame, row, 0, "Text Encoder Sequence Length", - tooltip="Number of tokens for captions") - components.entry(frame, row, 1, self.ui_state, "text_encoder_sequence_length") - row += 1 - - def __create_text_encoder_n_frame( - self, - master, - row: int, - i: int, - supports_include: bool = False, - supports_layer_skip: bool = True, - supports_sequence_length: bool = False, - ): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - row = 0 - - suffix = f"_{i}" if i > 1 else "" - - if supports_include: - # include text encoder - components.label(frame, row, 0, f"Include Text Encoder {i}", - tooltip=f"Includes text encoder {i} in the training run") - components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.include") - row += 1 - - # train text encoder - components.label(frame, row, 0, f"Train Text Encoder {i}", - tooltip=f"Enables training the text encoder {i} model") - components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.train") - row += 1 - - # train text encoder embedding - components.label(frame, row, 0, f"Train Text Encoder {i} Embedding", - tooltip=f"Enables training embeddings for the text encoder {i} model") - components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.train_embedding") - row += 1 - - # dropout - components.label(frame, row, 0, "Dropout Probability", - tooltip=f"The Probability for dropping the text encoder {i} conditioning") - components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}.dropout_probability") - row += 1 - - # train text encoder epochs - components.label(frame, row, 0, "Stop Training After", - tooltip=f"When to stop training the text encoder {i}") - components.time_entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}.stop_training_after", - f"text_encoder{suffix}.stop_training_after_unit", supports_time_units=False) - row += 1 - - # text encoder learning rate - components.label(frame, row, 0, f"Text Encoder {i} Learning Rate", - tooltip=f"The learning rate of the text encoder {i}. Overrides the base learning rate") - components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}.learning_rate") - row += 1 - - if supports_layer_skip: - # text encoder layer skip (clip skip) - components.label(frame, row, 0, f"Text Encoder {i} Clip Skip", - tooltip="The number of additional clip layers to skip. 0 = the model default") - components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}_layer_skip") - row += 1 - - if supports_sequence_length: - # text encoder sequence length - components.label(frame, row, 0, f"Text Encoder {i} Sequence Length", - tooltip="Overrides the number of tokens used for captions. If empty, the model default is used, which is 512 on Flux. Comfy samples with 256 tokens though. 77 is the default only for backwards compatibility.") - components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}_sequence_length") - row += 1 - - def __create_embedding_frame(self, master, row): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - - # embedding learning rate - components.label(frame, 0, 0, "Embeddings Learning Rate", - tooltip="The learning rate of embeddings. Overrides the base learning rate") - components.entry(frame, 0, 1, self.ui_state, "embedding_learning_rate") - - # preserve embedding norm - components.label(frame, 1, 0, "Preserve Embedding Norm", - tooltip="Rescales each trained embedding to the median embedding norm") - components.switch(frame, 1, 1, self.ui_state, "preserve_embedding_norm") - - def __create_unet_frame(self, master, row): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - - # train unet - components.label(frame, 0, 0, "Train UNet", - tooltip="Enables training the UNet model") - components.switch(frame, 0, 1, self.ui_state, "unet.train") - - # train unet epochs - components.label(frame, 1, 0, "Stop Training After", - tooltip="When to stop training the UNet") - components.time_entry(frame, 1, 1, self.ui_state, "unet.stop_training_after", "unet.stop_training_after_unit", - supports_time_units=False) - - # unet learning rate - components.label(frame, 2, 0, "UNet Learning Rate", - tooltip="The learning rate of the UNet. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "unet.learning_rate") - - # rescale noise scheduler to zero terminal SNR - rescale_label = components.label(frame, 3, 0, "Rescale Noise Scheduler + V-pred", - tooltip="Rescales the noise scheduler to a zero terminal signal to noise ratio and switches the model to a v-prediction target") - rescale_label.configure(wraplength=130, justify="left") - components.switch(frame, 3, 1, self.ui_state, "rescale_noise_scheduler_to_zero_terminal_snr") - - def __create_prior_frame(self, master, row): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - - # train prior - components.label(frame, 0, 0, "Train Prior", - tooltip="Enables training the Prior model") - components.switch(frame, 0, 1, self.ui_state, "prior.train") - - # train prior epochs - components.label(frame, 1, 0, "Stop Training After", - tooltip="When to stop training the Prior") - components.time_entry(frame, 1, 1, self.ui_state, "prior.stop_training_after", "prior.stop_training_after_unit", - supports_time_units=False) - - # prior learning rate - components.label(frame, 2, 0, "Prior Learning Rate", - tooltip="The learning rate of the Prior. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "prior.learning_rate") - - def __create_transformer_frame(self, master, row, supports_guidance_scale: bool = False, supports_force_attention_mask: bool = True): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - - # train transformer - components.label(frame, 0, 0, "Train Transformer", - tooltip="Enables training the Transformer model") - components.switch(frame, 0, 1, self.ui_state, "transformer.train") - - # train transformer epochs - components.label(frame, 1, 0, "Stop Training After", - tooltip="When to stop training the Transformer") - components.time_entry(frame, 1, 1, self.ui_state, "transformer.stop_training_after", "transformer.stop_training_after_unit", - supports_time_units=False) - - # transformer learning rate - components.label(frame, 2, 0, "Transformer Learning Rate", - tooltip="The learning rate of the Transformer. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "transformer.learning_rate") - - if supports_force_attention_mask: - # transformer learning rate - components.label(frame, 3, 0, "Force Attention Mask", - tooltip="Force enables passing of a text embedding attention mask to the transformer. This can improve training on shorter captions.") - components.switch(frame, 3, 1, self.ui_state, "transformer.attention_mask") - - if supports_guidance_scale: - # guidance scale - components.label(frame, 4, 0, "Guidance Scale", - tooltip="The guidance scale of guidance distilled models passed to the transformer during training.") - components.entry(frame, 4, 1, self.ui_state, "transformer.guidance_scale") - - def __create_noise_frame(self, master, row, supports_generalized_offset_noise: bool = False, supports_dynamic_timestep_shifting: bool = False): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - - # offset noise weight - components.label(frame, 0, 0, "Offset Noise Weight", - tooltip="The weight of offset noise added to each training step") - components.entry(frame, 0, 1, self.ui_state, "offset_noise_weight") - - if supports_generalized_offset_noise: - # generalized offset noise weight - generalised_offset_label = components.label(frame, 1, 0, "Generalized Offset Noise", - tooltip="Per-timestep 'brightness knob' instead of a fixed offset - steadier training, better starts, and improved very dark/bright images. Compatible with V-pred and Eps-pred. Start with 0.02 and adjust as needed.") - generalised_offset_label.configure(wraplength=130, justify="left") - components.switch(frame, 1, 1, self.ui_state, "generalized_offset_noise") - - # perturbation noise weight - components.label(frame, 2, 0, "Perturbation Noise Weight", - tooltip="The weight of perturbation noise added to each training step") - components.entry(frame, 2, 1, self.ui_state, "perturbation_noise_weight") - - # timestep distribution - components.label(frame, 3, 0, "Timestep Distribution", - tooltip="Selects the function to sample timesteps during training", - wide_tooltip=True) - components.options_adv(frame, 3, 1, [str(x) for x in list(TimestepDistribution)], self.ui_state, "timestep_distribution", - adv_command=self.__open_timestep_distribution_window) - - # min noising strength - components.label(frame, 4, 0, "Min Noising Strength", - tooltip="Specifies the minimum noising strength used during training. This can help to improve composition, but prevents finer details from being trained") - components.entry(frame, 4, 1, self.ui_state, "min_noising_strength", required=True) - - # max noising strength - components.label(frame, 5, 0, "Max Noising Strength", - tooltip="Specifies the maximum noising strength used during training. This can be useful to reduce overfitting, but also reduces the impact of training samples on the overall image composition") - components.entry(frame, 5, 1, self.ui_state, "max_noising_strength", required=True) - - # noising weight - components.label(frame, 6, 0, "Noising Weight", - tooltip="Controls the weight parameter of the timestep distribution function. Use the preview to see more details.") - components.entry(frame, 6, 1, self.ui_state, "noising_weight", required=True) - - # noising bias - components.label(frame, 7, 0, "Noising Bias", - tooltip="Controls the bias parameter of the timestep distribution function. Use the preview to see more details.") - components.entry(frame, 7, 1, self.ui_state, "noising_bias", required=True) - - # timestep shift - components.label(frame, 8, 0, "Timestep Shift", - tooltip="Shift the timestep distribution. Use the preview to see more details.") - components.entry(frame, 8, 1, self.ui_state, "timestep_shift", required=True) - - if supports_dynamic_timestep_shifting: - # dynamic timestep shifting - components.label(frame, 9, 0, "Dynamic Timestep Shifting", - tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) - components.switch(frame, 9, 1, self.ui_state, "dynamic_timestep_shifting") - - - - def __create_masked_frame(self, master, row): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - - # Masked Training - components.label(frame, 0, 0, "Masked Training", - tooltip="Masks the training samples to let the model focus on certain parts of the image. When enabled, one mask image is loaded for each training sample.") - components.switch(frame, 0, 1, self.ui_state, "masked_training") - - # unmasked probability - components.label(frame, 1, 0, "Unmasked Probability", - tooltip="When masked training is enabled, specifies the number of training steps done on unmasked samples") - components.entry(frame, 1, 1, self.ui_state, "unmasked_probability", - extra_validate=check_range(lower=0, upper=1, message="Unmasked probability must be between 0 and 1")) - - # unmasked weight - components.label(frame, 2, 0, "Unmasked Weight", - tooltip="When masked training is enabled, specifies the loss weight of areas outside the masked region") - components.entry(frame, 2, 1, self.ui_state, "unmasked_weight", - extra_validate=check_range(lower=0, upper=1, message="Unmasked weight must be between 0 and 1")) - - # normalize masked area loss - components.label(frame, 3, 0, "Normalize Masked Area Loss", - tooltip="When masked training is enabled, normalizes the loss for each sample based on the sizes of the masked region") - components.switch(frame, 3, 1, self.ui_state, "normalize_masked_area_loss") - - # masked prior preservation - components.label(frame, 4, 0, "Masked Prior Preservation Weight", - tooltip="Preserves regions outside the mask using the original untrained model output as a target. Only available for LoRA training. If enabled, use a low unmasked weight.") - components.entry(frame, 4, 1, self.ui_state, "masked_prior_preservation_weight", - extra_validate=check_range(lower=0, upper=1, message="Masked prior preservation weight must be between 0 and 1")) - - # use custom conditioning image - components.label(frame, 5, 0, "Custom Conditioning Image", - tooltip="When custom conditioning image is enabled, will use png postfix with -condlabel instead of automatically generated.It's suitable for special scenarios, such as object removal, allowing the model to learn a certain behavior concept") - components.switch(frame, 5, 1, self.ui_state, "custom_conditioning_image") - - def __create_loss_frame(self, master, row, supports_vb_loss: bool = False): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - - # MSE Strength - components.label(frame, 0, 0, "MSE Strength", - tooltip="Mean Squared Error strength for custom loss settings. Strengths should generally sum to 1.") - components.entry(frame, 0, 1, self.ui_state, "mse_strength", required=True) - - # MAE Strength - components.label(frame, 1, 0, "MAE Strength", - tooltip="Mean Absolute Error strength for custom loss settings. Strengths should generally sum to 1.") - components.entry(frame, 1, 1, self.ui_state, "mae_strength", required=True) - - # log-cosh Strength - components.label(frame, 2, 0, "log-cosh Strength", - tooltip="Log - Hyperbolic cosine Error strength for custom loss settings. Strengths should generally sum to 1.") - components.entry(frame, 2, 1, self.ui_state, "log_cosh_strength", required=True) - - # Huber Strength - components.label(frame, 3, 0, "Huber Strength", - tooltip="Huber loss strength for custom loss settings. Less sensitive to outliers than MSE. Strengths should generally sum to 1.") - components.entry(frame, 3, 1, self.ui_state, "huber_strength", required=True) - - # Huber Delta - components.label(frame, 4, 0, "Huber Delta", - tooltip="Delta parameter for huber loss") - components.entry(frame, 4, 1, self.ui_state, "huber_delta", required=True) - - if supports_vb_loss: - # VB Strength - components.label(frame, 5, 0, "VB Strength", - tooltip="Variational lower-bound strength for custom loss settings. Should be set to 1 for variational diffusion models") - components.entry(frame, 5, 1, self.ui_state, "vb_loss_strength", required=True) - - # Loss Weight function - components.label(frame, 6, 0, "Loss Weight Function", - tooltip="Choice of loss weight function. Can help the model learn details more accurately.") - components.options(frame, 6, 1, [str(x) for x in list(LossWeight) - if x.supports_flow_matching() == self.train_config.model_type.is_flow_matching() - or x == LossWeight.CONSTANT - ], - self.ui_state, "loss_weight_fn") - - row = 7 - - # Loss weight strength - if not self.train_config.model_type.is_flow_matching(): - components.label(frame, row, 0, "Gamma", - tooltip="Inverse strength of loss weighting. Range: 1-20, only applies to Min SNR and P2.") - components.entry(frame, row, 1, self.ui_state, "loss_weight_strength", - extra_validate=check_range(lower=1, upper=20, message="Gamma must be between 1 and 20")) - row += 1 - - # Loss Scaler - components.label(frame, row, 0, "Loss Scaler", - tooltip="Selects the type of loss scaling to use during training. Functionally equated as: Loss * selection") - components.options(frame, row, 1, [str(x) for x in list(LossScaler)], self.ui_state, "loss_scaler") - row += 1 - - def __create_layer_frame(self, master, row): - cls = create.get_model_setup_class(self.train_config.model_type, self.train_config.training_method) - presets = cls.LAYER_PRESETS if cls is not None else {"full": []} - components.layer_filter_entry(master, row, 0, self.ui_state, - preset_var_name="layer_filter_preset", presets=presets, - preset_label="Layer Filter", - preset_tooltip="Select a preset defining which layers to train, or select 'Custom' to define your own.\nA blank 'custom' field or 'Full' will train all layers.", - entry_var_name="layer_filter", - entry_tooltip="Comma-separated list of diffusion layers to train. Regular expressions (if toggled) are supported. Any model layer with a matching name will be trained", - regex_var_name="layer_filter_regex", - regex_tooltip="If enabled, layer filter patterns are interpreted as regular expressions. Otherwise, simple substring matching is used.", - ) - - - def __on_layer_filter_preset_change(self): - if not self.layer_selector: + def _restore_scheduler_config(self, variable): + if not hasattr(self, 'lr_scheduler_adv_comp'): return - selected = self.ui_state.get_var("layer_filter_preset").get() - self.__preset_set_layer_choice(selected) - - def __hide_layer_entry(self): - if self.layer_entry and self.layer_entry.winfo_manager(): - self.layer_entry.grid_remove() - - def __show_layer_entry(self): - if self.layer_entry and not self.layer_entry.winfo_manager(): - self.layer_entry.grid() - - def __open_optimizer_params_window(self): - window = OptimizerParamsWindow(self.master, self.train_config, self.ui_state) - self.master.wait_window(window) + state = "normal" if self.controller.is_custom_scheduler_value(variable) else "disabled" + self.lr_scheduler_adv_comp.configure(state=state) - def __open_scheduler_params_window(self): - window = SchedulerParamsWindow(self.master, self.train_config, self.ui_state) - self.master.wait_window(window) + def _open_optimizer_params_window(self): + self.master.wait_window(self.controller.open_optimizer_params_window(self.master, self.ui_state, CtkOptimizerParamsWindowView)) - def __open_timestep_distribution_window(self): - window = TimestepDistributionWindow(self.master, self.train_config, self.ui_state) - self.master.wait_window(window) + def _open_scheduler_params_window(self): + self.master.wait_window(self.controller.open_scheduler_params_window(self.master, self.ui_state, CtkSchedulerParamsWindowView)) - def __open_offloading_window(self): - window = OffloadingWindow(self.master, self.train_config, self.ui_state) - self.master.wait_window(window) - - def __restore_optimizer_config(self, *args): - optimizer_config = change_optimizer(self.train_config) - self.ui_state.get_var("optimizer").update(optimizer_config) - - def __restore_scheduler_config(self, variable): - if not hasattr(self, 'lr_scheduler_adv_comp'): - return + def _open_timestep_distribution_window(self): + self.master.wait_window(self.controller.open_timestep_distribution_window(self.master, self.ui_state, CtkTimestepDistributionWindowView)) - if variable == "CUSTOM": - self.lr_scheduler_adv_comp.configure(state="normal") - else: - self.lr_scheduler_adv_comp.configure(state="disabled") + def _open_offloading_window(self): + self.master.wait_window(self.controller.open_offloading_window(self.master, self.ui_state, CtkOffloadingWindowView)) diff --git a/modules/ui/CtkVideoToolUIView.py b/modules/ui/CtkVideoToolUIView.py index c3291e6ea..c272c891c 100644 --- a/modules/ui/CtkVideoToolUIView.py +++ b/modules/ui/CtkVideoToolUIView.py @@ -1,33 +1,23 @@ -import concurrent.futures -import math -import os -import pathlib -import random -import shlex -import subprocess -import threading -import webbrowser -from fractions import Fraction from tkinter import filedialog +from modules.ui.BaseVideoToolUIView import BaseVideoToolUIView +from modules.ui.VideoToolUIController import VideoToolUIController from modules.util.image_util import load_image -from modules.util.path_util import SUPPORTED_VIDEO_EXTENSIONS -from modules.util.ui import components +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState -import av import customtkinter as ctk -import cv2 -import scenedetect -from PIL import Image +PAD = ctk_components.PAD -class VideoToolUI(ctk.CTkToplevel): - def __init__( - self, - parent, - *args, **kwargs, - ): + +class CtkVideoToolUIView(BaseVideoToolUIView, ctk.CTkToplevel): + def __init__(self, parent, controller: VideoToolUIController, *args, **kwargs): ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseVideoToolUIView.__init__(self, ctk_components) + + self.controller = controller + ui_state = CtkUIState(self, controller.args) self.title("Video Tools") self.geometry("600x720") @@ -42,836 +32,97 @@ def __init__( tabview = ctk.CTkTabview(self) tabview.grid(row=0, column=0, sticky="nsew") - self.clip_extract_tab = self.__clip_extract_tab(tabview.add("extract clips")) - self.image_extract_tab = self.__image_extract_tab(tabview.add("extract images")) - self.video_download_tab = self.__video_download_tab(tabview.add("download")) - self.status_bar(self) - - def status_bar(self, master): + clip_frame = ctk.CTkScrollableFrame(tabview.add("extract clips"), fg_color="transparent") + clip_frame.grid_columnconfigure(0, weight=0, minsize=120) + clip_frame.grid_columnconfigure(1, weight=0, minsize=200) + clip_frame.grid_columnconfigure(2, weight=0) + clip_frame.grid_columnconfigure(3, weight=1) + self.build_clip_extract_tab(clip_frame, controller, ui_state) + clip_frame.pack(fill="both", expand=1) + + image_frame = ctk.CTkScrollableFrame(tabview.add("extract images"), fg_color="transparent") + image_frame.grid_columnconfigure(0, weight=0, minsize=120) + image_frame.grid_columnconfigure(1, weight=0, minsize=200) + image_frame.grid_columnconfigure(2, weight=0) + image_frame.grid_columnconfigure(3, weight=1) + self.build_image_extract_tab(image_frame, controller, ui_state) + image_frame.pack(fill="both", expand=1) + + download_frame = ctk.CTkScrollableFrame(tabview.add("download"), fg_color="transparent") + download_frame.grid_columnconfigure(0, weight=0, minsize=120) + download_frame.grid_columnconfigure(1, weight=0, minsize=200) + download_frame.grid_columnconfigure(2, weight=0) + download_frame.grid_columnconfigure(3, weight=1) + self.build_video_download_tab(download_frame, controller, ui_state) + download_frame.pack(fill="both", expand=1) + + self._build_status_bar(self) + + def _build_status_bar(self, master): frame = ctk.CTkFrame(master, fg_color="transparent") frame.grid(row=1, column=0) frame.grid_columnconfigure(0, weight=0, minsize=160) frame.grid_columnconfigure(1, weight=0, minsize=300) frame.grid_columnconfigure(2, weight=1) - #create preview image preview_path = "resources/icons/icon.png" preview = load_image(preview_path, 'RGB') preview.thumbnail((150, 150)) - self.preview_image= ctk.CTkImage(light_image=preview, size=preview.size) + self.preview_image = ctk.CTkImage(light_image=preview, size=preview.size) self.preview_image_label = ctk.CTkLabel( master=frame, text="Preview image", image=self.preview_image, height=150, width=150, compound="top") self.preview_image_label.grid(row=0, column=0, sticky="nw", padx=5, pady=5) - #displays progress and messages that also go to terminal self.status_label = ctk.CTkTextbox(master=frame, width=400, height=160, wrap="word", border_width=2) self.status_label.insert(index="1.0", text="Current status") self.status_label.configure(state="disabled") self.status_label.grid(row=0, column=1, sticky="ne", padx=5, pady=5) - def __clip_extract_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=120) - frame.grid_columnconfigure(1, weight=0, minsize=200) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # single video - components.label(frame, 0, 0, "Single Video", - tooltip="Link to single video file to process.") - self.clip_single_entry = ctk.CTkEntry(frame, width=190) - self.clip_single_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) - self.clip_single_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_file(self.clip_single_entry, - [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))] - )) - self.clip_single_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 0, 2, "Extract Single", - command=lambda: self.__extract_clips_button(False)) - - # time range - components.label(frame, 1, 0, " Time Range", - tooltip="Time range to limit selection for single video, \ - format as hour:minute:second, minute:second, or seconds.") - self.clip_time_start_entry = ctk.CTkEntry(frame, width=100) - self.clip_time_start_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.clip_time_start_entry.insert(0, "00:00:00") - self.clip_time_end_entry = ctk.CTkEntry(frame, width=100) - self.clip_time_end_entry.grid(row=1, column=1, sticky="e", padx=5, pady=5) - self.clip_time_end_entry.insert(0, "99:99:99") - - # directory of videos - components.label(frame, 2, 0, "Directory", - tooltip="Path to directory with multiple videos to process, including in subdirectories.") - self.clip_list_entry = ctk.CTkEntry(frame, width=190) - self.clip_list_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - self.clip_list_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.clip_list_entry)) - self.clip_list_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 2, 2, "Extract Directory", - command=lambda: self.__extract_clips_button(True)) - - # output directory - components.label(frame, 3, 0, "Output", - tooltip="Path to folder where extracted clips will be saved.") - self.clip_output_entry = ctk.CTkEntry(frame, width=190) - self.clip_output_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) - self.clip_output_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.clip_output_entry)) - self.clip_output_button.grid(row=3, column=1, sticky="e", padx=5, pady=5) - - # output to subdirectories - self.output_subdir_clip = ctk.BooleanVar(self, False) - components.label(frame, 4, 0, "Output to\nSubdirectories", - tooltip="If enabled, files are saved to subfolders based on filename and input directory. \ - Otherwise will all be saved to the top level of the output directory.") - self.output_subdir_clip_entry = ctk.CTkSwitch(frame, variable=self.output_subdir_clip, text="") - self.output_subdir_clip_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) - - # split at cuts - self.split_at_cuts = ctk.BooleanVar(self, False) - components.label(frame, 5, 0, "Split at Cuts", - tooltip="If enabled, detect cuts in the input video and split at those points. \ - Otherwise will split at any point, and clips may contain cuts.") - self.split_cuts_entry = ctk.CTkSwitch(frame, variable=self.split_at_cuts, text="") - self.split_cuts_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) - - # maximum length - components.label(frame, 6, 0, "Max Length (s)", - tooltip="Maximum length in seconds for saved clips, larger clips will be broken into multiple small clips.") - self.clip_length_entry = ctk.CTkEntry(frame, width=220) - self.clip_length_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) - self.clip_length_entry.insert(0, "3") - - # Set FPS - components.label(frame, 7, 0, "Set FPS", - tooltip="FPS to convert output videos to, set to 0 to keep original rate.") - self.clip_fps_entry = ctk.CTkEntry(frame, width=220) - self.clip_fps_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) - self.clip_fps_entry.insert(0, "24.0") - - # Remove borders - self.clip_bordercrop = ctk.BooleanVar(self, False) - components.label(frame, 8, 0, "Remove Borders", - tooltip="Remove black borders from output clip") - self.clip_bordercrop_entry = ctk.CTkSwitch(frame, variable=self.clip_bordercrop, text="") - self.clip_bordercrop_entry.grid(row=8, column=1, sticky="w", padx=5, pady=5) - - # Crop Variation - components.label(frame, 9, 0, "Crop Variation", - tooltip="Output clips will be randomly cropped to +- the base aspect ratio, \ - somewhat biased towards making square videos. Set to 0 to use only base aspect.") - self.clip_crop_entry = ctk.CTkEntry(frame, width=220) - self.clip_crop_entry.grid(row=9, column=1, sticky="w", padx=5, pady=5) - self.clip_crop_entry.insert(0, "0.2") - - # object filter - currently unused, may implement in future - # components.label(frame, 9, 0, "Object Filter", - # tooltip="Detect general features using Haar-Cascade classifier, and choose how to deal with clips where it is detected") - # components.options(frame, 9, 1, ["NONE", "FACE", "EYE", "BODY"], self.video_ui_state, "filter_object") - # components.options(frame, 9, 2, ["INCLUDE", "EXCLUDE", "SUBFOLDER"], self.video_ui_state, "filter_behavior") - - frame.pack(fill="both", expand=1) - return frame - - def __image_extract_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=120) - frame.grid_columnconfigure(1, weight=0, minsize=200) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # single video - components.label(frame, 0, 0, "Single Video", - tooltip="Link to single video file to process.") - self.image_single_entry = ctk.CTkEntry(frame, width=190) - self.image_single_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) - self.image_single_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_file(self.image_single_entry, - [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))] - )) - self.image_single_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 0, 2, "Extract Single", - command=lambda: self.__extract_images_button(False)) - - # time range - components.label(frame, 1, 0, " Time Range", - tooltip="Time range to limit selection for single video, \ - format as hour:minute:second, minute:second, or seconds.") - self.image_time_start_entry = ctk.CTkEntry(frame, width=100) - self.image_time_start_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.image_time_start_entry.insert(0, "00:00:00") - self.image_time_end_entry = ctk.CTkEntry(frame, width=100) - self.image_time_end_entry.grid(row=1, column=1, sticky="e", padx=5, pady=5) - self.image_time_end_entry.insert(0, "99:99:99") - - # directory of videos - components.label(frame, 2, 0, "Directory", - tooltip="Path to directory with multiple videos to process, including in subdirectories.") - self.image_list_entry = ctk.CTkEntry(frame, width=190) - self.image_list_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - self.image_list_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.image_list_entry)) - self.image_list_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 2, 2, "Extract Directory", - command=lambda: self.__extract_images_button(True)) - - # output directory - components.label(frame, 3, 0, "Output", - tooltip="Path to folder where extracted images will be saved.") - self.image_output_entry = ctk.CTkEntry(frame, width=190) - self.image_output_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) - self.image_output_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.image_output_entry)) - self.image_output_button.grid(row=3, column=1, sticky="e", padx=5, pady=5) - - # output to subdirectories - self.output_subdir_img = ctk.BooleanVar(self, False) - components.label(frame, 4, 0, "Output to\nSubdirectories", - tooltip="If enabled, files are saved to subfolders based on filename and input directory. \ - Otherwise will all be saved to the top level of the output directory.") - self.output_subdir_img_entry = ctk.CTkSwitch(frame, variable=self.output_subdir_img, text="") - self.output_subdir_img_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) - - # image capture rate - components.label(frame, 5, 0, "Images/sec", - tooltip="Number of images to capture per second of video. \ - Images will be taken at semi-random frames around the specified frequency.") - self.capture_rate_entry = ctk.CTkEntry(frame, width=220) - self.capture_rate_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) - self.capture_rate_entry.insert(0, "0.5") - - # blur removal - components.label(frame, 6, 0, "Blur Removal", - tooltip="Threshold for removal of blurry images, relative to all others. \ - For example at 0.2, the blurriest 20%% of the final selected frames will not be saved.") - self.blur_threshold_entry = ctk.CTkEntry(frame, width=220) - self.blur_threshold_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) - self.blur_threshold_entry.insert(0, "0.2") - - # Remove borders - self.image_bordercrop = ctk.BooleanVar(self, False) - components.label(frame, 7, 0, "Remove Borders", - tooltip="Remove black borders from output image") - self.image_bordercrop_entry = ctk.CTkSwitch(frame, variable=self.image_bordercrop, text="") - self.image_bordercrop_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) - - # Crop Variation - components.label(frame, 8, 0, "Crop Variation", - tooltip="Output images will be randomly cropped to +- the base aspect ratio, \ - somewhat biased towards making square images. Set to 0 to use only base sapect.") - self.image_crop_entry = ctk.CTkEntry(frame, width=220) - self.image_crop_entry.grid(row=8, column=1, sticky="w", padx=5, pady=5) - self.image_crop_entry.insert(0, "0.2") - - # # object filter - currently unused, may implement in future - # components.label(frame, 5, 0, "Object Filter", - # tooltip="Detect general features using Haar-Cascade classifier, and choose how to deal with clips where it is detected") - # components.options(frame, 5, 1, ["NONE", "FACE", "EYE", "BODY"], self.video_ui_state, "filter_object") - # components.options(frame, 5, 2, ["INCLUDE", "EXCLUDE", "SUBFOLDER"], self.video_ui_state, "filter_behavior") - - frame.pack(fill="both", expand=1) - return frame - - def __video_download_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=120) - frame.grid_columnconfigure(1, weight=0, minsize=200) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # link - components.label(frame, 0, 0, "Single Link", - tooltip="Link to video/playlist to download. Uses yt-dlp, supports youtube, twitch, instagram, and many other sites.") - self.download_link_entry = ctk.CTkEntry(frame, width=220) - self.download_link_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) - components.button(frame, 0, 2, "Download Link", command=lambda: self.__download_button(False)) - - # link list - components.label(frame, 1, 0, "Link List", - tooltip="Path to txt file with list of links separated by newlines.") - self.download_list_entry = ctk.CTkEntry(frame, width=190) - self.download_list_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.download_list_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_file(self.download_list_entry, [("Text file", ".txt")])) - self.download_list_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 1, 2, "Download List", command=lambda: self.__download_button(True)) - - # output directory - components.label(frame, 2, 0, "Output", - tooltip="Path to folder where downloaded videos will be saved.") - self.download_output_entry = ctk.CTkEntry(frame, width=190) - self.download_output_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - self.download_output_button = ctk.CTkButton(frame, width=30, text="...", command=lambda: self.__browse_for_dir(self.download_output_entry)) - self.download_output_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) - - # additional args - 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.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 --format mp4") - components.button(frame, 3, 2, "yt-dlp info", - command=lambda: webbrowser.open("https://github.com/yt-dlp/yt-dlp?tab=readme-ov-file#usage-and-options", new=0, autoraise=False)) - - frame.pack(fill="both", expand=1) - return frame - - def __browse_for_dir(self, entry_box): - # get the path from the user - path = filedialog.askdirectory() - # set the path to the entry box - # delete entry box text - entry_box.focus_set() - entry_box.delete(0, ctk.END) - entry_box.insert(0, path) - self.focus_set() - - def __browse_for_file(self, entry_box, filetypes): - # get the path from the user - path = filedialog.askopenfilename(filetypes=filetypes) - # set the path to the entry box - # delete entry box text - entry_box.focus_set() - entry_box.delete(0, ctk.END) - entry_box.insert(0, path) - self.focus_set() - - def __get_vid_paths(self, batch_mode: bool, input_path_single: str, input_path_dir: str): - input_videos = [] - if not batch_mode: - path = pathlib.Path(input_path_single) - if path.is_file(): - vid = cv2.VideoCapture(str(path)) - ok = False - try: - if vid.isOpened(): - ok, _ = vid.read() - finally: - vid.release() - if ok: - return [path] - else: - self.__update_status("Invalid video file!") - return [] - else: - self.__update_status("No file specified, or invalid file path!") - return [] - else: - input_videos = [] - if not pathlib.Path(input_path_dir).is_dir() or input_path_dir == "": - self.__update_status("Invalid input directory!") - return [] - # Only traverse supported extensions to avoid opening every file. - lower_exts = {e.lower() for e in SUPPORTED_VIDEO_EXTENSIONS} - for path in pathlib.Path(input_path_dir).rglob("*"): - if path.is_file() and path.suffix.lower() in lower_exts: - vid = cv2.VideoCapture(str(path)) - ok = False - try: - if vid.isOpened(): - ok, _ = vid.read() - finally: - vid.release() - if ok: - input_videos.append(path) - self.__update_status(f'Found {len(input_videos)} videos to process') - return input_videos - - def __run_in_thread(self, target, *args): - """Clear status box and run target function in a daemon thread.""" + def _create_textbox(self, master, row, col, width, height, ui_state, var_name): + var = ui_state.get_var(var_name) + textbox = ctk.CTkTextbox(master, width=width, height=height, border_width=2) + textbox.insert("1.0", var.get()) + textbox.grid(row=row, column=col, rowspan=2, sticky="w", padx=PAD, pady=PAD) + + def on_text_change(event=None): + var.set(textbox.get("1.0", "end-1c")) + + textbox.bind("", on_text_change) + return textbox + + def _create_browse_dir_button(self, master, row, ui_state, var_name): + def browse(): + path = filedialog.askdirectory() + if path: + ui_state.get_var(var_name).set(path) + self.focus_set() + + button = ctk.CTkButton(master, width=30, text="...", command=browse) + button.grid(row=row, column=1, sticky="e", padx=PAD, pady=PAD) + return button + + def _create_browse_file_button(self, master, row, ui_state, var_name, filetypes): + def browse(): + path = filedialog.askopenfilename(filetypes=filetypes) + if path: + ui_state.get_var(var_name).set(path) + self.focus_set() + + button = ctk.CTkButton(master, width=30, text="...", command=browse) + button.grid(row=row, column=1, sticky="e", padx=PAD, pady=PAD) + return button + + def update_status(self, status_text: str): self.status_label.configure(state="normal") - self.status_label.delete(index1="1.0", index2="end") + self.status_label.insert(index="end", text=status_text + "\n") self.status_label.configure(state="disabled") - t = threading.Thread(target=target, args=args) - t.daemon = True - t.start() - - @staticmethod - def __parse_timestamp_to_frames(timestamp: str, fps: float) -> int: - return int(sum(int(x) * 60 ** i for i, x in enumerate(reversed(timestamp.split(':')))) * fps) - - def __get_safe_fps(self, video: cv2.VideoCapture, video_path: str) -> float: - fps = video.get(cv2.CAP_PROP_FPS) or 0.0 - if fps <= 0: - self.__update_status(f'Warning: Could not read FPS for "{os.path.basename(video_path)}". Falling back to 30 FPS.') - return 30.0 - return fps - - @staticmethod - def __get_output_dir(use_subdir: bool, batch_mode: bool, output_entry: str, - video_path, input_dir: str) -> str: - if use_subdir and batch_mode: - return os.path.join(output_entry, - os.path.splitext(os.path.relpath(video_path, input_dir))[0]) - elif use_subdir: - return os.path.join(output_entry, - os.path.splitext(os.path.basename(video_path))[0]) - return output_entry - - def __get_random_aspect(self, height: int, width: int, variation: float) -> tuple[int, int, int, int]: - # Return original dimensions and no offset if variation is zero - if variation == 0: - return 0, height, 0, width - - old_aspect = height/width - variation_scaled = old_aspect*variation - if old_aspect > 1.2: #tall image - new_aspect = min(4.0, max(1.0, random.triangular(old_aspect-(variation_scaled*1.5), old_aspect+(variation_scaled/2), old_aspect))) - elif old_aspect < 0.85: #wide image - new_aspect = max(0.25, min(1.0, random.triangular(old_aspect-(variation_scaled/2), old_aspect+(variation_scaled*1.5), old_aspect))) - else: #square image - new_aspect = random.triangular(old_aspect-variation_scaled, old_aspect+variation_scaled) - - new_aspect = round(new_aspect, 2) - #keep the height the same if reducing width, and vice versa - if new_aspect > old_aspect: - new_height = int(height) - new_width = int(width*(old_aspect/new_aspect)) - elif new_aspect < old_aspect: - new_height = int(height*(new_aspect/old_aspect)) - new_width = int(width) - else: - new_height = int(height) - new_width = int(width) - - #random offset in dimension that was cropped - position_x = random.randint(0, width-new_width) - position_y = random.randint(0, height-new_height) - return position_y, new_height, position_x, new_width - - def find_main_contour(self, frame): - #outline image to find main content and exclude black bars often present on letterboxed videos - frame_grayscale = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) - _, frame_thresh = cv2.threshold(frame_grayscale, 15, 255, cv2.THRESH_BINARY) - frame_contours, _ = cv2.findContours(frame_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) - if frame_contours: - #select largest contour by area - frame_maincontour = max(frame_contours, key=lambda c: cv2.contourArea(c)) - x1, y1, w1, h1 = cv2.boundingRect(frame_maincontour) - else: #fallback if no contours detected - x1 = 0 - y1 = 0 - h1, w1, _ = frame.shape - - #if bounding box did not detect the correct area, likely due to all-black frame - if not frame_contours or h1 < 10 or w1 < 10: - x1 = 0 - y1 = 0 - h1, w1, _ = frame.shape - return x1, y1, w1, h1 - - def __extract_clips_button(self, batch_mode: bool): - self.__run_in_thread(self.__extract_clips_multi, batch_mode) - - def __extract_clips_multi(self, batch_mode: bool): - if not pathlib.Path(self.clip_output_entry.get()).is_dir() or self.clip_output_entry.get() == "": - self.__update_status("Invalid output directory!") - return - - # validate numeric inputs - try: - max_length = float(self.clip_length_entry.get()) - crop_variation = float(self.clip_crop_entry.get()) - target_fps = float(self.clip_fps_entry.get()) - input_single_entry = self.clip_single_entry.get() - input_multiple_entry = self.clip_list_entry.get() - output_entry = self.clip_output_entry.get() - except ValueError: - self.__update_status("Invalid numeric input for Max Length, Crop Variation, or FPS.") - return - if max_length <= 0.25: - self.__update_status("Max Length of clips must be > 0.25 seconds.") - return - if target_fps < 0: - self.__update_status("Target FPS must be a positive number (or 0 to skip fps re-encoding).") - return - if not (0.0 <= crop_variation < 1.0): - self.__update_status("Crop Variation must be between 0.0 and 1.0.") - return - input_videos = self.__get_vid_paths(batch_mode, input_single_entry, input_multiple_entry) - if len(input_videos) == 0: # exit if no paths found - return - - with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: - for video_path in input_videos: - output_directory = self.__get_output_dir( - self.output_subdir_clip_entry.get(), batch_mode, - output_entry, video_path, input_multiple_entry) - time_start = "00:00:00" if batch_mode else str(self.clip_time_start_entry.get()) - time_end = "99:99:99" if batch_mode else str(self.clip_time_end_entry.get()) - executor.submit(self.__extract_clips, - str(video_path), time_start, time_end, max_length, - self.split_at_cuts.get(), bool(self.clip_bordercrop_entry.get()), - crop_variation, target_fps, output_directory) - - if batch_mode: - self.__update_status(f'Clip extraction from all videos in "{input_multiple_entry}" complete') - else: - self.__update_status(f'Clip extraction from "{input_single_entry}" complete') - - def __extract_clips(self, video_path: str, timestamp_min: str, timestamp_max: str, max_length: float, - split_at_cuts: bool, remove_borders: bool, crop_variation: float, target_fps: float, output_dir: str): - video = cv2.VideoCapture(video_path) - vid_fps = self.__get_safe_fps(video, video_path) - max_length_frames = int(max_length * vid_fps) - min_length_frames = max(int(0.25 * vid_fps), 1) - total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) or 0 - timestamp_max_frame = min(self.__parse_timestamp_to_frames(timestamp_max, vid_fps), max(total_frames - 1, 0)) - timestamp_min_frame = min(self.__parse_timestamp_to_frames(timestamp_min, vid_fps), timestamp_max_frame) - - if split_at_cuts: - #use scenedetect to find cuts, based on start/end frame number - self.__update_status(f'Detecting scenes in "{os.path.basename(video_path)}"') - timecode_list = scenedetect.detect( - video_path=str(video_path), - detector=scenedetect.AdaptiveDetector(), - start_time=int(timestamp_min_frame), - end_time=int(timestamp_max_frame)) - scene_list = [(x[0].get_frames(), x[1].get_frames()) for x in timecode_list] - if not scene_list: - scene_list = [(timestamp_min_frame, timestamp_max_frame)] - else: - scene_list = [(timestamp_min_frame, timestamp_max_frame)] - - scene_list_split = [] - for scene in scene_list: - length = scene[1]-scene[0] - if length > max_length_frames: #check for any scenes longer than max length - n = math.ceil(length/max_length_frames) #divide into n new scenes - new_length = int(length/n) - new_splits = range(scene[0], scene[1]+min_length_frames, new_length) #divide clip into closest chunks to max_length - for i, _n in enumerate(new_splits[:-1]): - if new_splits[i + 1] - new_splits[i] > min_length_frames: - scene_list_split.append((new_splits[i], new_splits[i + 1])) - elif length > (min_length_frames + 2): - # Trim first/last frame to avoid transition artifacts - scene_list_split.append((scene[0] + 1, scene[1] - 1)) - - self.__update_status(f'Video "{os.path.basename(video_path)}" being split into {len(scene_list_split)} clips in "{output_dir}"') - - with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: - futures = [ - executor.submit(self.__save_clip, scene, video_path, target_fps, - remove_borders, crop_variation, output_dir) - for scene in scene_list_split - ] - for future in concurrent.futures.as_completed(futures): - exc = future.exception() - if exc is not None: - self.__update_status(f'Error saving clip: {exc}') - - video.release() - - def __save_clip(self, scene: tuple[int, int], video_path: str, target_fps: float, - remove_borders: bool, crop_variation: float, output_dir: str): - basename, ext = os.path.splitext(os.path.basename(video_path)) - video = cv2.VideoCapture(str(video_path)) - fps = self.__get_safe_fps(video, video_path) - os.makedirs(output_dir, exist_ok=True) - output_name = f'{output_dir}{os.sep}{basename}_{scene[0]}-{scene[1]}' - output_ext = ".mp4" - - video.set(cv2.CAP_PROP_POS_FRAMES, (scene[1] + scene[0])//2) #set to middle of scene - frame_number = int(video.get(cv2.CAP_PROP_POS_FRAMES)) - success, frame = video.read() - if not success or frame is None: - self.__update_status(f'Failed to read frame from "{os.path.basename(video_path)}" at {int(frame_number)}. Skipping clip.') - video.release() - return - - # Blend random frames to detect borders, avoiding incorrect crop from black frames - if remove_borders: - frame_blend = frame - for i in range(5): - random_frame = random.randint(scene[0], scene[1]) - video.set(cv2.CAP_PROP_POS_FRAMES, random_frame) - success, frame = video.read() - if not success or frame is None: - continue - a = 1/(i+1) - b = 1-a - frame_blend = cv2.addWeighted(frame, a, frame_blend, b, 0) - x1, y1, w1, h1 = self.find_main_contour(frame_blend) - else: - x1 = 0 - y1 = 0 - h1, w1, _ = frame.shape - - y2, h2, x2, w2 = self.__get_random_aspect(h1, w1, crop_variation) - # Ensure dimensions are even, required - h2 -= h2 % 2 - w2 -= w2 % 2 - print(end='\x1b[2K') #clear terminal so next line can overwrite it - print(f'Saving frames {scene[0]}-{scene[1]} at size {w2}x{h2}', end="\r") - video.set(cv2.CAP_PROP_POS_FRAMES, (scene[1] + scene[0])//2) - success, frame = video.read() - if success: - try: - preview = Image.fromarray( - cv2.cvtColor(frame[y1+y2:y1+y2+h2, x1+x2:x1+x2+w2], cv2.COLOR_BGR2RGB)) - preview.thumbnail((150, 150)) - self.preview_image.configure(light_image=preview, size=preview.size) - #truncate filename of long files so UI doesn't shift around - filename_truncated = basename + ext if len(basename) < 20 else basename[:18] + ".." + ext - self.preview_image_label.configure( - text=f'{filename_truncated}\nFrames: {scene[0]}-{scene[1]}\nSize: {w2}x{h2}') - except Exception: - pass - video.release() - - if target_fps <= 0: - target_fps = fps - - output_path = f'{output_name}{output_ext}' - self.__write_clip_av(video_path, output_path, scene, fps, target_fps, - x1 + x2, y1 + y2, w2, h2) - - @staticmethod - def __write_clip_av(video_path: str, output_path: str, scene: tuple[int, int], - src_fps: float, target_fps: float, - crop_x: int, crop_y: int, crop_w: int, crop_h: int): - start_sec = scene[0] / src_fps - end_sec = scene[1] / src_fps - rate_frac = Fraction(target_fps).limit_denominator(10000) - stream_time_base = Fraction(rate_frac.denominator, rate_frac.numerator) - - with av.open(video_path) as input_container: - in_video = input_container.streams.video[0] - in_video.thread_type = 'AUTO' - in_audio = input_container.streams.audio[0] if input_container.streams.audio else None - - with av.open(output_path, mode='w') as output_container: - out_video = output_container.add_stream('libx264', rate=rate_frac) - out_video.width = crop_w - out_video.height = crop_h - out_video.pix_fmt = 'yuv420p' - out_video.time_base = stream_time_base - - out_audio = output_container.add_stream_from_template(in_audio) if in_audio else None - - input_container.seek(int(start_sec * 1_000_000)) - - out_frame_idx = 0 - out_time_step = 1.0 / target_fps - video_done = False - decode_streams = [s for s in (in_video, in_audio) if s is not None] - - for packet in input_container.demux(decode_streams): - if packet.stream == in_video: - if video_done: - continue - for frame in packet.decode(): - if frame.time is None or frame.time < start_sec: - continue - if frame.time >= end_sec: - video_done = True - break - - # FPS conversion: skip frames when source fps > target fps - if frame.time < start_sec + out_frame_idx * out_time_step: - continue - - img = frame.to_ndarray(format='bgr24') - cropped = img[crop_y:crop_y + crop_h, crop_x:crop_x + crop_w] - out_frame = av.VideoFrame.from_ndarray(cropped, format='bgr24') - out_frame.pts = out_frame_idx - out_frame.time_base = stream_time_base - - for out_pkt in out_video.encode(out_frame): - output_container.mux(out_pkt) - out_frame_idx += 1 - - elif packet.stream == in_audio and out_audio is not None: - if packet.dts is None: - continue - pkt_time = float(packet.pts * packet.time_base) - if pkt_time < start_sec or pkt_time >= end_sec: - continue - # Re-timestamp audio relative to clip start - packet.pts = int((pkt_time - start_sec) / packet.time_base) - packet.dts = packet.pts - packet.stream = out_audio - output_container.mux(packet) - - # Flush video encoder - for pkt in out_video.encode(): - output_container.mux(pkt) - - def __extract_images_button(self, batch_mode: bool): - self.__run_in_thread(self.__extract_images_multi, batch_mode) - - def __extract_images_multi(self, batch_mode : bool): - if not pathlib.Path(self.image_output_entry.get()).is_dir() or self.image_output_entry.get() == "": - self.__update_status("Invalid output directory!") - return - - # validate numeric inputs - try: - capture_rate = float(self.capture_rate_entry.get()) - blur_threshold = float(self.blur_threshold_entry.get()) - crop_variation = float(self.image_crop_entry.get()) - input_single_entry = self.image_single_entry.get() - input_multiple_entry = self.image_list_entry.get() - output_entry = self.image_output_entry.get() - except ValueError: - self.__update_status("Invalid numeric input for Images/sec, Blur Removal, or Crop Variation.") - return - if capture_rate <= 0: - self.__update_status("Images/sec must be > 0.") - return - if not (0.0 <= blur_threshold < 1.0): - self.__update_status("Blur Removal must be between 0.0 and 1.0.") - return - if not (0.0 <= crop_variation < 1.0): - self.__update_status("Crop Variation must be between 0.0 and 1.0.") - return - - input_videos = self.__get_vid_paths(batch_mode, input_single_entry, input_multiple_entry) - if not input_videos: - return - - with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: - for video_path in input_videos: - output_directory = self.__get_output_dir( - self.output_subdir_img_entry.get(), batch_mode, - output_entry, video_path, input_multiple_entry) - time_start = "00:00:00" if batch_mode else str(self.image_time_start_entry.get()) - time_end = "99:99:99" if batch_mode else str(self.image_time_end_entry.get()) - executor.submit(self.__save_frames, - str(video_path), time_start, time_end, capture_rate, - blur_threshold, self.image_bordercrop.get(), - crop_variation, output_directory) - if batch_mode: - self.__update_status(f'Image extraction from all videos in {input_multiple_entry} complete') - else: - self.__update_status(f'Image extraction from "{input_single_entry}" complete') - - def __save_frames(self, video_path: str, timestamp_min: str, timestamp_max: str, capture_rate: float, - blur_threshold: float, remove_borders: bool, crop_variation: float, output_dir: str): - video = cv2.VideoCapture(video_path) - vid_fps = self.__get_safe_fps(video, video_path) - if capture_rate <= 0: - self.__update_status("Images/sec must be > 0.") - video.release() - return - image_rate = max(int(vid_fps / capture_rate), 1) - total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) or 0 - timestamp_max_frame = min(self.__parse_timestamp_to_frames(timestamp_max, vid_fps), max(total_frames - 1, 0)) - timestamp_min_frame = min(self.__parse_timestamp_to_frames(timestamp_min, vid_fps), timestamp_max_frame) - frame_range = range(timestamp_min_frame, timestamp_max_frame, image_rate) - frame_list = [] - - for n in frame_range: - #pick frame from random triangular distribution around center of each "chunk" of the video - frame = abs(int(random.triangular(n-(image_rate/2), n+(image_rate/2)))) - frame = max(0, min(frame, max(total_frames - 1, 0))) - frame_list.append(frame) - - self.__update_status(f'Video "{os.path.basename(video_path)}" will be split into {len(frame_list)} images in "{output_dir}"') - - output_list = [] - for f in frame_list: - video.set(cv2.CAP_PROP_POS_FRAMES, f) - success, frame = video.read() - if success and frame is not None: - frame_grayscale = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) - frame_sharpness = cv2.Laplacian(frame_grayscale, cv2.CV_64F).var() - output_list.append((f, frame_sharpness)) - - if not output_list: - self.__update_status(f'No frames extracted from "{os.path.basename(video_path)}" in the selected range.') - video.release() - return - - output_list_sorted = sorted(output_list, key=lambda x: x[1]) - cutoff = int(blur_threshold * len(output_list_sorted)) - output_list_cut = output_list_sorted[cutoff:] - self.__update_status(f'{cutoff} blurriest images have been dropped from "{os.path.basename(video_path)}"') - - basename, ext = os.path.splitext(os.path.basename(video_path)) - os.makedirs(output_dir, exist_ok=True) - - for f in output_list_cut: - filename = f'{output_dir}{os.sep}{basename}_{f[0]}.jpg' - video.set(cv2.CAP_PROP_POS_FRAMES, f[0]) - success, frame = video.read() - - #crop out borders of frame - if remove_borders and success and frame is not None: - x1, y1, w1, h1 = self.find_main_contour(frame) - frame_cropped = frame[y1:y1+h1, x1:x1+w1] - else: - frame_cropped = frame if success and frame is not None else None - if frame_cropped is not None: - x1 = 0 - y1 = 0 - h1, w1, _ = frame_cropped.shape - - y2, h2, x2, w2 = self.__get_random_aspect(h1, w1, crop_variation) - - if success and frame is not None and frame_cropped is not None: - print(end='\x1b[2K') #clear terminal so next line can overwrite it - print(f'Saving frame {f[0]} at size {w2}x{h2}', end="\r") - try: - preview = Image.fromarray( - cv2.cvtColor(frame_cropped[y2:y2+h2, x2:x2+w2], cv2.COLOR_BGR2RGB)) - preview.thumbnail((150, 150)) - filename_truncated = basename + ext if len(basename) < 20 else basename[:17] + "..." + ext - self.preview_image.configure(light_image=preview, size=preview.size) - self.preview_image_label.configure(text=f'{filename_truncated}\nFrame: {f[0]}\nSize: {w2}x{h2}') - except Exception: - pass # preview update is non-critical - - cv2.imwrite(filename, frame_cropped[y2:y2+h2, x2:x2+w2]) - video.release() - - def __download_button(self, batch_mode: bool): - self.__run_in_thread(self.__download_multi, batch_mode) - - def __update_status(self, status_text: str): - print(status_text) + def clear_status(self): self.status_label.configure(state="normal") - self.status_label.insert(index="end", text=status_text + "\n") + self.status_label.delete(index1="1.0", index2="end") self.status_label.configure(state="disabled") - def __download_multi(self, batch_mode: bool): - if not pathlib.Path(self.download_output_entry.get()).is_dir() or self.download_output_entry.get() == "": - self.__update_status("Invalid output directory!") - return - - if not batch_mode: - ydl_urls = [self.download_link_entry.get()] - elif batch_mode: - ydl_path = pathlib.Path(self.download_list_entry.get()) - if ydl_path.is_file() and ydl_path.suffix.lower() == ".txt": - with open(ydl_path) as file: - ydl_urls = file.readlines() - else: - self.__update_status("Invalid link list!") - return - - with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: - for url in ydl_urls: - executor.submit(self.__download_video, - url.strip(), self.download_output_entry.get(), - self.download_args_entry.get("0.0", ctk.END)) - - self.__update_status(f'Completed {len(ydl_urls)} downloads.') - - def __download_video(self, url: str, output_dir: str, output_args: str): - url = (url or "").strip() - if not url: - self.__update_status("Empty URL, skipping download.") - return - - #Respect quotes and split into list to run as yt-dlp command - additional_args = shlex.split(output_args.strip()) if output_args and output_args.strip() else [] - cmd = ["yt-dlp", "-o", "%(title)s.%(ext)s", "-P", output_dir] + additional_args + [url] - - self.__update_status(f'Downloading {url}') - subprocess.run(cmd) - self.__update_status(f'Download {url} done!') + def update_preview(self, preview_image, label_text: str): + self.preview_image.configure(light_image=preview_image, size=preview_image.size) + self.preview_image_label.configure(text=label_text) diff --git a/modules/ui/GenerateCaptionsWindowController.py b/modules/ui/GenerateCaptionsWindowController.py index 1690879f1..2b2411b28 100644 --- a/modules/ui/GenerateCaptionsWindowController.py +++ b/modules/ui/GenerateCaptionsWindowController.py @@ -1,133 +1,28 @@ -import contextlib -import tkinter as tk -from tkinter import filedialog - -from modules.util.ui.ui_utils import set_window_icon - -import customtkinter as ctk - - -class GenerateCaptionsWindow(ctk.CTkToplevel): - def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs): - """ - Window for generating captions for a folder of images - - Parameters: - parent (`Tk`): the parent window - path (`str`): the path to the folder - parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox - """ - super().__init__(parent, *args, **kwargs) +class GenerateCaptionsWindowController: + def __init__(self, parent): self.parent = parent + self.view = None - if path is None: - path = "" - - self.mode_var = ctk.StringVar(self, "Create if absent") - self.modes = ["Replace all captions", "Create if absent", "Add as new line"] - self.model_var = ctk.StringVar(self, "Blip") - self.models = ["Blip", "Blip2", "WD14 VIT v2"] - - self.title("Batch generate captions") - self.geometry("360x360") - self.resizable(True, True) - - self.frame = ctk.CTkFrame(self, width=600, height=300) - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - - self.model_label = ctk.CTkLabel(self.frame, text="Model", width=100) - self.model_label.grid(row=0, column=0, sticky="w", padx=5, pady=5) - self.model_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.model_var, values=self.models, dynamic_resizing=False, width=200) - self.model_dropdown.grid(row=0, column=1, sticky="w", padx=5, pady=5) - - self.path_label = ctk.CTkLabel(self.frame, text="Folder", width=100) - self.path_label.grid(row=1, column=0, sticky="w",padx=5, pady=5) - self.path_entry = ctk.CTkEntry(self.frame, width=150) - self.path_entry.insert(0, path) - self.path_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.path_button = ctk.CTkButton(self.frame, width=30, text="...", command=lambda: self.browse_for_path(self.path_entry)) - self.path_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) - - self.caption_label = ctk.CTkLabel(self.frame, text="Initial Caption", width=100) - self.caption_label.grid(row=2, column=0, sticky="w", padx=5, pady=5) - self.caption_entry = ctk.CTkEntry(self.frame, width=200) - self.caption_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - - self.prefix_label = ctk.CTkLabel(self.frame, text="Caption Prefix", width=100) - self.prefix_label.grid(row=3, column=0, sticky="w", padx=5, pady=5) - self.prefix_entry = ctk.CTkEntry(self.frame, width=200) - self.prefix_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) + def create_window(self, parent_window, path, parent_include_subdirectories, view_cls): + self.view = view_cls(parent_window, self, path, parent_include_subdirectories) + return self.view - self.postfix_label = ctk.CTkLabel(self.frame, text="Caption Postfix", width=100) - self.postfix_label.grid(row=4, column=0, sticky="w", padx=5, pady=5) - self.postfix_entry = ctk.CTkEntry(self.frame, width=200) - self.postfix_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) - - self.mode_label = ctk.CTkLabel(self.frame, text="Mode", width=100) - self.mode_label.grid(row=5, column=0, sticky="w", padx=5, pady=5) - self.mode_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.mode_var, values=self.modes, dynamic_resizing=False, width=200) - self.mode_dropdown.grid(row=5, column=1, sticky="w", padx=5, pady=5) - - self.include_subdirectories_label = ctk.CTkLabel(self.frame, text="Include subfolders", width=100) - self.include_subdirectories_label.grid(row=6, column=0, sticky="w", padx=5, pady=5) - self.include_subdirectories_var = ctk.BooleanVar(self, parent_include_subdirectories) - self.include_subdirectories_switch = ctk.CTkSwitch(self.frame, text="", variable=self.include_subdirectories_var) - self.include_subdirectories_switch.grid(row=6, column=1, sticky="w", padx=5, pady=5) - - self.progress_label = ctk.CTkLabel(self.frame, text="Progress: 0/0", width=100) - self.progress_label.grid(row=7, column=0, sticky="w", padx=5, pady=5) - self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) - self.progress.grid(row=7, column=1, sticky="w", padx=5, pady=5) - - self.create_captions_button = ctk.CTkButton(self.frame, text="Create Captions", width=310, command=self.create_captions) - self.create_captions_button.grid(row=8, column=0, columnspan=2, sticky="w", padx=5, pady=5) - - self.frame.pack(fill="both", expand=True) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def browse_for_path(self, entry_box): - # get the path from the user - path = filedialog.askdirectory() - # set the path to the entry box - # delete entry box text - entry_box.focus_set() - entry_box.delete(0, filedialog.END) - entry_box.insert(0, path) - self.focus_set() - - def set_progress(self, value, max_value): - progress = value / max_value - self.progress.set(progress) - self.progress_label.configure(text=f"{value}/{max_value}") - self.progress.update() - - def create_captions(self): - self.parent.load_captioning_model(self.model_var.get()) + def create_captions(self, model_name, path, initial_caption, caption_prefix, caption_postfix, mode_str, include_subdirectories): + self.parent.load_captioning_model(model_name) mode = { "Replace all captions": "replace", "Create if absent": "fill", "Add as new line": "add", - }[self.mode_var.get()] + }[mode_str] self.parent.captioning_model.caption_folder( - sample_dir=self.path_entry.get(), - initial_caption=self.caption_entry.get(), - caption_prefix=self.prefix_entry.get(), - caption_postfix=self.postfix_entry.get(), + sample_dir=path, + initial_caption=initial_caption, + caption_prefix=caption_prefix, + caption_postfix=caption_postfix, mode=mode, - progress_callback=self.set_progress, - include_subdirectories=self.include_subdirectories_var.get(), + progress_callback=self.view.set_progress, + include_subdirectories=include_subdirectories, ) self.parent.load_image() - - def destroy(self): - with contextlib.suppress(tk.TclError): - self.grab_release() - - super().destroy() diff --git a/modules/ui/GenerateMasksWindowController.py b/modules/ui/GenerateMasksWindowController.py index daff0d3d5..6c154ef46 100644 --- a/modules/ui/GenerateMasksWindowController.py +++ b/modules/ui/GenerateMasksWindowController.py @@ -1,127 +1,14 @@ -import contextlib -import tkinter as tk -from tkinter import filedialog - -from modules.util.ui.ui_utils import set_window_icon - -import customtkinter as ctk - - -class GenerateMasksWindow(ctk.CTkToplevel): - def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs): - """ - Window for generating masks for a folder of images - - Parameters: - parent (`Tk`): the parent window - path (`str`): the path to the folder - parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox - """ - super().__init__(parent, *args, **kwargs) - +class GenerateMasksWindowController: + def __init__(self, parent): self.parent = parent - if path is None: - path = "" - - self.mode_var = ctk.StringVar(self, "Create if absent") - self.modes = ["Replace all masks", "Create if absent", "Add to existing", "Subtract from existing", "Blend with existing"] - self.model_var = ctk.StringVar(self, "ClipSeg") - self.models = ["ClipSeg", "Rembg", "Rembg-Human", "Hex Color"] - - self.title("Batch generate masks") - self.geometry("360x430") - self.resizable(True, True) - - self.frame = ctk.CTkFrame(self, width=600, height=300) - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - - self.model_label = ctk.CTkLabel(self.frame, text="Model", width=100) - self.model_label.grid(row=0, column=0, sticky="w", padx=5, pady=5) - self.model_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.model_var, values=self.models, dynamic_resizing=False, width=200) - self.model_dropdown.grid(row=0, column=1, sticky="w", padx=5, pady=5) - - self.path_label = ctk.CTkLabel(self.frame, text="Folder", width=100) - self.path_label.grid(row=1, column=0, sticky="w",padx=5, pady=5) - self.path_entry = ctk.CTkEntry(self.frame, width=150) - self.path_entry.insert(0, path) - self.path_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.path_button = ctk.CTkButton(self.frame, width=30, text="...", command=lambda: self.browse_for_path(self.path_entry)) - self.path_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) - - self.prompt_label = ctk.CTkLabel(self.frame, text="Prompt", width=100) - self.prompt_label.grid(row=2, column=0, sticky="w",padx=5, pady=5) - self.prompt_entry = ctk.CTkEntry(self.frame, width=200) - self.prompt_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - - self.mode_label = ctk.CTkLabel(self.frame, text="Mode", width=100) - self.mode_label.grid(row=3, column=0, sticky="w", padx=5, pady=5) - self.mode_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.mode_var, values=self.modes, dynamic_resizing=False, width=200) - self.mode_dropdown.grid(row=3, column=1, sticky="w", padx=5, pady=5) + self.view = None - self.threshold_label = ctk.CTkLabel(self.frame, text="Threshold", width=100) - self.threshold_label.grid(row=4, column=0, sticky="w", padx=5, pady=5) - self.threshold_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="0.0 - 1.0") - self.threshold_entry.insert(0, "0.3") - self.threshold_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + def create_window(self, parent_window, path, parent_include_subdirectories, view_cls): + self.view = view_cls(parent_window, self, path, parent_include_subdirectories) + return self.view - self.smooth_label = ctk.CTkLabel(self.frame, text="Smooth", width=100) - self.smooth_label.grid(row=5, column=0, sticky="w", padx=5, pady=5) - self.smooth_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="5") - self.smooth_entry.insert(0, 5) - self.smooth_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) - - self.expand_label = ctk.CTkLabel(self.frame, text="Expand", width=100) - self.expand_label.grid(row=6, column=0, sticky="w", padx=5, pady=5) - self.expand_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="10") - self.expand_entry.insert(0, 10) - self.expand_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) - - self.alpha_label = ctk.CTkLabel(self.frame, text="Alpha", width=100) - self.alpha_label.grid(row=7, column=0, sticky="w", padx=5, pady=5) - self.alpha_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="1") - self.alpha_entry.insert(0, 1) - self.alpha_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) - - self.include_subdirectories_label = ctk.CTkLabel(self.frame, text="Include subfolders", width=100) - self.include_subdirectories_label.grid(row=8, column=0, sticky="w", padx=5, pady=5) - self.include_subdirectories_var = ctk.BooleanVar(self, parent_include_subdirectories) - self.include_subdirectories_switch = ctk.CTkSwitch(self.frame, text="", variable=self.include_subdirectories_var) - self.include_subdirectories_switch.grid(row=8, column=1, sticky="w", padx=5, pady=5) - - self.progress_label = ctk.CTkLabel(self.frame, text="Progress: 0/0", width=100) - self.progress_label.grid(row=9, column=0, sticky="w", padx=5, pady=5) - self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) - self.progress.grid(row=9, column=1, sticky="w", padx=5, pady=5) - - self.create_masks_button = ctk.CTkButton(self.frame, text="Create Masks", width=310, command=self.create_masks) - self.create_masks_button.grid(row=10, column=0, columnspan=2, sticky="w", padx=5, pady=5) - - self.frame.pack(fill="both", expand=True) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def browse_for_path(self, entry_box): - # get the path from the user - path = filedialog.askdirectory() - # set the path to the entry box - # delete entry box text - entry_box.focus_set() - entry_box.delete(0, filedialog.END) - entry_box.insert(0, path) - self.focus_set() - - def set_progress(self, value, max_value): - progress = value / max_value - self.progress.set(progress) - self.progress_label.configure(text=f"{value}/{max_value}") - self.progress.update() - - def create_masks(self): - self.parent.load_masking_model(self.model_var.get()) + def create_masks(self, model_name, path, prompt, mode_str, alpha_str, threshold_str, smooth_str, expand_str, include_subdirectories): + self.parent.load_masking_model(model_name) mode = { "Replace all masks": "replace", @@ -129,23 +16,17 @@ def create_masks(self): "Add to existing": "add", "Subtract from existing": "subtract", "Blend with existing": "blend", - }[self.mode_var.get()] + }[mode_str] self.parent.masking_model.mask_folder( - sample_dir=self.path_entry.get(), - prompts=[self.prompt_entry.get()], + sample_dir=path, + prompts=[prompt], mode=mode, - alpha=float(self.alpha_entry.get()), - threshold=float(self.threshold_entry.get()), - smooth_pixels=int(self.smooth_entry.get()), - expand_pixels=int(self.expand_entry.get()), - progress_callback=self.set_progress, - include_subdirectories=self.include_subdirectories_var.get(), + alpha=float(alpha_str), + threshold=float(threshold_str), + smooth_pixels=int(smooth_str), + expand_pixels=int(expand_str), + progress_callback=self.view.set_progress, + include_subdirectories=include_subdirectories, ) self.parent.load_image() - - def destroy(self): - with contextlib.suppress(tk.TclError): - self.grab_release() - - super().destroy() diff --git a/modules/ui/LoraTabController.py b/modules/ui/LoraTabController.py new file mode 100644 index 000000000..46aa4aab3 --- /dev/null +++ b/modules/ui/LoraTabController.py @@ -0,0 +1,22 @@ + +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.DataType import DataType +from modules.util.enum.ModelType import PeftType + + +class LoraTabController: + def __init__(self, config: TrainConfig): + self.train_config = config + + def get_peft_types(self) -> list[tuple[str, PeftType]]: + return [ + ("LoRA", PeftType.LORA), + ("LoHa", PeftType.LOHA), + ("OFT v2", PeftType.OFT_2), + ] + + def get_lora_weight_dtypes(self) -> list[tuple[str, DataType]]: + return [ + ("float32", DataType.FLOAT_32), + ("bfloat16", DataType.BFLOAT_16), + ] diff --git a/modules/ui/ModelTabController.py b/modules/ui/ModelTabController.py new file mode 100644 index 000000000..69f603f78 --- /dev/null +++ b/modules/ui/ModelTabController.py @@ -0,0 +1,13 @@ + + +from modules.util import create +from modules.util.config.TrainConfig import TrainConfig + + +class ModelTabController: + def __init__(self, config: TrainConfig): + self.train_config = config + + def get_presets(self) -> dict: + cls = create.get_model_setup_class(self.train_config.model_type, self.train_config.training_method) + return cls.LAYER_PRESETS if cls is not None else {"full": []} diff --git a/modules/ui/MuonAdamWindowController.py b/modules/ui/MuonAdamWindowController.py new file mode 100644 index 000000000..834cae400 --- /dev/null +++ b/modules/ui/MuonAdamWindowController.py @@ -0,0 +1,28 @@ +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.Optimizer import Optimizer +from modules.util.optimizer_util import OPTIMIZER_DEFAULT_PARAMETERS + +MUON_AUX_ADAM_DEFAULTS = { + "beta1": 0.9, + "beta2": 0.999, + "eps": 1e-8, + "weight_decay": 0.0, +} + + + + +class MuonAdamWindowController: + def __init__(self, config: TrainConfig, parent_optimizer_type: Optimizer): + self.config = config + self.parent_optimizer_type = parent_optimizer_type + + def get_title(self) -> str: + if self.parent_optimizer_type == Optimizer.MUON: + return "Muon's Auxiliary AdamW Settings" + return "Muon_adv's Auxiliary AdamW_adv Settings" + + def get_adam_params_def(self) -> dict: + if self.parent_optimizer_type == Optimizer.MUON: + return MUON_AUX_ADAM_DEFAULTS + return OPTIMIZER_DEFAULT_PARAMETERS[Optimizer.ADAMW_ADV] diff --git a/modules/ui/OffloadingWindowController.py b/modules/ui/OffloadingWindowController.py new file mode 100644 index 000000000..cc9bad142 --- /dev/null +++ b/modules/ui/OffloadingWindowController.py @@ -0,0 +1,6 @@ +from modules.util.config.TrainConfig import TrainConfig + + +class OffloadingWindowController: + def __init__(self, config: TrainConfig): + self.config = config diff --git a/modules/ui/OptimizerParamsWindowController.py b/modules/ui/OptimizerParamsWindowController.py index 16063c26c..37f838907 100644 --- a/modules/ui/OptimizerParamsWindowController.py +++ b/modules/ui/OptimizerParamsWindowController.py @@ -1,7 +1,5 @@ -import contextlib -from tkinter import TclError -from modules.ui.MuonAdamWindow import MUON_AUX_ADAM_DEFAULTS, MuonAdamWindow +from modules.ui.MuonAdamWindowController import MUON_AUX_ADAM_DEFAULTS from modules.util.config.TrainConfig import TrainConfig, TrainOptimizerConfig from modules.util.enum.Optimizer import Optimizer from modules.util.optimizer_util import ( @@ -10,261 +8,27 @@ load_optimizer_defaults, update_optimizer_config, ) -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState -import customtkinter as ctk +class OptimizerParamsWindowController: + def __init__(self, config: TrainConfig): + self.config = config -class OptimizerParamsWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - train_config: TrainConfig, - ui_state, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) + def restore_optimizer_config(self, ui_state): + optimizer_config = change_optimizer(self.config) + ui_state.get_var("optimizer").update(optimizer_config) - self.parent = parent - self.train_config = train_config - self.ui_state = ui_state - self.optimizer_ui_state = ui_state.get_var("optimizer") - self.protocol("WM_DELETE_WINDOW", self.on_window_close) - self.muon_adam_button = None + def load_defaults(self, ui_state): + optimizer_config = load_optimizer_defaults(self.config) + ui_state.get_var("optimizer").update(optimizer_config) - self.title("Optimizer Settings") - self.geometry("800x500") - self.resizable(True, True) - - self.grid_rowconfigure(0, weight=1) - self.grid_rowconfigure(1, weight=0) - self.grid_columnconfigure(0, weight=1) - - self.frame = ctk.CTkScrollableFrame(self, fg_color="transparent") - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) - self.frame.grid_columnconfigure(2, minsize=50) - self.frame.grid_columnconfigure(3, weight=0) - self.frame.grid_columnconfigure(4, weight=1) - - components.button(self, 1, 0, "ok", command=self.on_window_close) - self.main_frame(self.frame) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def main_frame(self, master): - # Optimizer - components.label(master, 0, 0, "Optimizer", - tooltip="The type of optimizer") - - # Create the optimizer dropdown menu and set the command - components.options(master, 0, 1, [str(x) for x in list(Optimizer)], self.optimizer_ui_state, "optimizer", - command=self.on_optimizer_change) - - # Defaults Button - components.label(master, 0, 3, "Optimizer Defaults", - tooltip="Load default settings for the selected optimizer") - components.button(self.frame, 0, 4, "Load Defaults", self.load_defaults, - tooltip="Load default settings for the selected optimizer") - - self.create_dynamic_ui(master) - - def clear_dynamic_ui(self, master): - with contextlib.suppress(TclError): - for widget in master.winfo_children(): - grid_info = widget.grid_info() - if int(grid_info["row"]) >= 1: - widget.destroy() - - def create_dynamic_ui( - self, - master, - ): - - # Lookup for the title and tooltip for a key - # @formatter:off - KEY_DETAIL_MAP = { - 'adam_w_mode': {'title': 'Adam W Mode', 'tooltip': 'Whether to use weight decay correction for Adam optimizer.', 'type': 'bool'}, - 'alpha': {'title': 'Alpha', 'tooltip': 'Smoothing parameter for RMSprop and others.', 'type': 'float'}, - 'amsgrad': {'title': 'AMSGrad', 'tooltip': 'Whether to use the AMSGrad variant for Adam.', 'type': 'bool'}, - 'beta1': {'title': 'Beta1', 'tooltip': 'optimizer_momentum term.', 'type': 'float'}, - 'beta2': {'title': 'Beta2', 'tooltip': 'Coefficients for computing running averages of gradient.', 'type': 'float'}, - 'beta3': {'title': 'Beta3', 'tooltip': 'Coefficient for computing the Prodigy stepsize.', 'type': 'float'}, - 'bias_correction': {'title': 'Bias Correction', 'tooltip': 'Whether to use bias correction in optimization algorithms like Adam.', 'type': 'bool'}, - 'block_wise': {'title': 'Block Wise', 'tooltip': 'Whether to perform block-wise model update.', 'type': 'bool'}, - 'capturable': {'title': 'Capturable', 'tooltip': 'Whether some property of the optimizer can be captured.', 'type': 'bool'}, - 'centered': {'title': 'Centered', 'tooltip': 'Whether to center the gradient before scaling. Great for stabilizing the training process.', 'type': 'bool'}, - 'clip_threshold': {'title': 'Clip Threshold', 'tooltip': 'Clipping value for gradients.', 'type': 'float'}, - 'd0': {'title': 'Initial D', 'tooltip': 'Initial D estimate for D-adaptation.', 'type': 'float'}, - 'd_coef': {'title': 'D Coefficient', 'tooltip': 'Coefficient in the expression for the estimate of d.', 'type': 'float'}, - 'dampening': {'title': 'Dampening', 'tooltip': 'Dampening for optimizer_momentum.', 'type': 'float'}, - 'decay_rate': {'title': 'Decay Rate', 'tooltip': 'Rate of decay for moment estimation.', 'type': 'float'}, - 'decouple': {'title': 'Decouple', 'tooltip': 'Use AdamW style optimizer_decoupled weight decay.', 'type': 'bool'}, - 'differentiable': {'title': 'Differentiable', 'tooltip': 'Whether the optimization function is optimizer_differentiable.', 'type': 'bool'}, - 'eps': {'title': 'EPS', 'tooltip': 'A small value to prevent division by zero.', 'type': 'float'}, - 'eps2': {'title': 'EPS 2', 'tooltip': 'A small value to prevent division by zero.', 'type': 'float'}, - 'foreach': {'title': 'ForEach', 'tooltip': 'Whether to use a foreach implementation if available. This implementation is usually faster.', 'type': 'bool'}, - 'fsdp_in_use': {'title': 'FSDP in Use', 'tooltip': 'Flag for using sharded parameters.', 'type': 'bool'}, - 'fused': {'title': 'Fused', 'tooltip': 'Whether to use a fused implementation if available. This implementation is usually faster and requires less memory.', 'type': 'bool'}, - 'fused_back_pass': {'title': 'Fused Back Pass', 'tooltip': 'Whether to fuse the back propagation pass with the optimizer step. This reduces VRAM usage, but is not compatible with gradient accumulation.', 'type': 'bool'}, - 'growth_rate': {'title': 'Growth Rate', 'tooltip': 'Limit for D estimate growth rate.', 'type': 'float'}, - 'initial_accumulator_value': {'title': 'Initial Accumulator Value', 'tooltip': 'Initial value for Adagrad optimizer.', 'type': 'float'}, - 'initial_accumulator': {'title': 'Initial Accumulator', 'tooltip': 'Sets the starting value for both moment estimates to ensure numerical stability and balanced adaptive updates early in training.', 'type': 'float'}, - 'is_paged': {'title': 'Is Paged', 'tooltip': 'Whether the optimizer\'s internal state should be paged to CPU.', 'type': 'bool'}, - 'log_every': {'title': 'Log Every', 'tooltip': 'Intervals at which logging should occur.', 'type': 'int'}, - 'lr_decay': {'title': 'LR Decay', 'tooltip': 'Rate at which learning rate decreases.', 'type': 'float'}, - 'max_unorm': {'title': 'Max Unorm', 'tooltip': 'Maximum value for gradient clipping by norms.', 'type': 'float'}, - 'maximize': {'title': 'Maximize', 'tooltip': 'Whether to optimizer_maximize the optimization function.', 'type': 'bool'}, - 'min_8bit_size': {'title': 'Min 8bit Size', 'tooltip': 'Minimum tensor size for 8-bit quantization.', 'type': 'int'}, - 'quant_block_size': {'title': 'Quant Block Size', 'tooltip': 'Size of a block of normalized 8-bit quantization data. Larger values increase memory efficiency at the cost of data precision.', 'type': 'int'}, - 'momentum': {'title': 'optimizer_momentum', 'tooltip': 'Factor to accelerate SGD in relevant direction.', 'type': 'float'}, - 'nesterov': {'title': 'Nesterov', 'tooltip': 'Whether to enable Nesterov optimizer_momentum.', 'type': 'bool'}, - 'no_prox': {'title': 'No Prox', 'tooltip': 'Whether to use proximity updates or not.', 'type': 'bool'}, - 'optim_bits': {'title': 'Optim Bits', 'tooltip': 'Number of bits used for optimization.', 'type': 'int'}, - 'percentile_clipping': {'title': 'Percentile Clipping', 'tooltip': 'Gradient clipping based on percentile values.', 'type': 'int'}, - 'relative_step': {'title': 'Relative Step', 'tooltip': 'Whether to use a relative step size.', 'type': 'bool'}, - 'safeguard_warmup': {'title': 'Safeguard Warmup', 'tooltip': 'Avoid issues during warm-up stage.', 'type': 'bool'}, - 'scale_parameter': {'title': 'Scale Parameter', 'tooltip': 'Whether to scale the parameter or not.', 'type': 'bool'}, - 'stochastic_rounding': {'title': 'Stochastic Rounding', 'tooltip': 'Stochastic rounding for weight updates. Improves quality when using bfloat16 weights.', 'type': 'bool'}, - 'use_bias_correction': {'title': 'Bias Correction', 'tooltip': 'Turn on Adam\'s bias correction.', 'type': 'bool'}, - 'use_triton': {'title': 'Use Triton', 'tooltip': 'Whether Triton optimization should be used.', 'type': 'bool'}, - 'warmup_init': {'title': 'Warmup Initialization', 'tooltip': 'Whether to warm-up the optimizer initialization.', 'type': 'bool'}, - 'weight_decay': {'title': 'Weight Decay', 'tooltip': 'Regularization to prevent overfitting.', 'type': 'float'}, - 'weight_lr_power': {'title': 'Weight LR Power', 'tooltip': 'During warmup, the weights in the average will be equal to lr raised to this power. Set to 0 for no weighting.', 'type': 'float'}, - 'decoupled_decay': {'title': 'Decoupled Decay', 'tooltip': 'If set as True, then the optimizer uses decoupled weight decay as in AdamW.', 'type': 'bool'}, - 'fixed_decay': {'title': 'Fixed Decay', 'tooltip': '(When Decoupled Decay is True:) Applies fixed weight decay when True; scales decay with learning rate when False.', 'type': 'bool'}, - 'rectify': {'title': 'Rectify', 'tooltip': 'Perform the rectified update similar to RAdam.', 'type': 'bool'}, - 'degenerated_to_sgd': {'title': 'Degenerated to SGD', 'tooltip': 'Performs SGD update when gradient variance is high.', 'type': 'bool'}, - 'k': {'title': 'K', 'tooltip': 'Number of vector projected per iteration.', 'type': 'int'}, - 'xi': {'title': 'Xi', 'tooltip': 'Term used in vector projections to avoid division by zero.', 'type': 'float'}, - 'n_sma_threshold': {'title': 'N SMA Threshold', 'tooltip': 'Number of SMA threshold.', 'type': 'int'}, - 'ams_bound': {'title': 'AMS Bound', 'tooltip': 'Whether to use the AMSBound variant.', 'type': 'bool'}, - 'r': {'title': 'R', 'tooltip': 'EMA factor.', 'type': 'float'}, - 'adanorm': {'title': 'AdaNorm', 'tooltip': 'Whether to use the AdaNorm variant', 'type': 'bool'}, - 'adam_debias': {'title': 'Adam Debias', 'tooltip': 'Only correct the denominator to avoid inflating step sizes early in training.', 'type': 'bool'}, - 'slice_p': {'title': 'Slice parameters', 'tooltip': 'Reduce memory usage by calculating LR adaptation statistics on only every pth entry of each tensor. For values greater than 1 this is an approximation to standard Prodigy. Values ~11 are reasonable.', 'type': 'int'}, - 'cautious': {'title': 'Cautious', 'tooltip': 'Whether to use the Cautious variant', 'type': 'bool'}, - 'weight_decay_by_lr': {'title': 'weight_decay_by_lr', 'tooltip': 'Automatically adjust weight decay based on lr', 'type': 'bool'}, - 'prodigy_steps': {'title': 'prodigy_steps', 'tooltip': 'Turn off Prodigy after N steps', 'type': 'int'}, - 'use_speed': {'title': 'use_speed', 'tooltip': 'use_speed method', 'type': 'bool'}, - 'split_groups': {'title': 'split_groups', 'tooltip': 'Use split groups when training multiple params(uNet,TE..)', 'type': 'bool'}, - 'split_groups_mean': {'title': 'split_groups_mean', 'tooltip': 'Use mean for split groups', 'type': 'bool'}, - 'factored': {'title': 'factored', 'tooltip': 'Use factored', 'type': 'bool'}, - 'factored_fp32': {'title': 'factored_fp32', 'tooltip': 'Use factored_fp32', 'type': 'bool'}, - 'use_stableadamw': {'title': 'use_stableadamw', 'tooltip': 'Use use_stableadamw for gradient scaling', 'type': 'bool'}, - 'use_cautious': {'title': 'use_cautious', 'tooltip': 'Use cautious method', 'type': 'bool'}, - 'use_grams': {'title': 'use_grams', 'tooltip': 'Use grams method', 'type': 'bool'}, - 'use_adopt': {'title': 'use_adopt', 'tooltip': 'Use adopt method', 'type': 'bool'}, - 'd_limiter': {'title': 'd_limiter', 'tooltip': 'Prevent over-estimated LRs when gradients and EMA are still stabilizing', 'type': 'bool'}, - 'use_schedulefree': {'title': 'use_schedulefree', 'tooltip': 'Use Schedulefree method', 'type': 'bool'}, - 'use_orthograd': {'title': 'use_orthograd', 'tooltip': 'Use orthograd method', 'type': 'bool'}, - 'nnmf_factor': {'title': 'Factored Optimizer', 'tooltip': 'Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. This is an approximation which may slightly alter training dynamics.', 'type': 'bool'}, - 'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization.', 'type': 'bool'}, - 'use_atan2': {'title': 'Atan2 Scaling', 'tooltip': 'A robust replacement for eps, which also incorporates gradient clipping, bounding and stabilizing the optimizer updates.', 'type': 'bool'}, - 'use_AdEMAMix': {'title': 'AdEMAMix EMA', 'tooltip': 'Adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, and accelerate the training.', 'type': 'bool'}, - 'beta3_ema': {'title': 'Beta3 EMA', 'tooltip': 'Coefficient for slow-moving EMA of AdEMAMix.', 'type': 'float'}, - 'beta1_warmup': {'title': 'Beta1 Warmup Steps', 'tooltip': 'Number of warmup steps to gradually increase beta1 from Minimum Beta1 Value to its final value. During warmup, beta1 increases linearly. leave it empty to disable warmup and use constant beta1.', 'type': 'int'}, - 'min_beta1': {'title': 'Minimum Beta1', 'tooltip': 'Starting beta1 value for warmup scheduling. Used only when beta1 warmup is enabled. Lower values allow faster initial adaptation, while higher values provide more smoothing. The final beta1 value is specified in the beta1 parameter.', 'type': 'float'}, - 'Simplified_AdEMAMix': {'title': 'Simplified AdEMAMix', 'tooltip': "Enables a simplified, single-EMA variant of AdEMAMix. Instead of blending two moving averages (fast and slow momentum), this version combines the raw current gradient (controlled by 'Grad α') directly with a single theory-based momentum. This makes the optimizer highly responsive to recent gradient information, which can accelerate training in all batch size scenarios when tuned correctly.", 'type': 'bool'}, - 'alpha_grad': {'title': 'Grad α', 'tooltip': 'Controls the mixing coefficient between raw gradients and momentum gradients in Simplified AdEMAMix. Higher values (e.g., 10-100) emphasize recent gradients, suitable for small batch sizes to reduce noise. Lower values (e.g., 0-1) emphasize historical gradients, suitable for large batch sizes for stability. Setting to 0 uses only momentum gradients without raw gradient contribution.', 'type': 'float'}, - 'kourkoutas_beta': {'title': 'Kourkoutas Beta', 'tooltip': 'Enables a layer-wise dynamic β₂ adaptation. This feature makes the optimizer more responsive to "spiky" gradients by lowering β₂ during periods of high variance, and more stable during calm periods by raising β₂ towards its maximum. It can significantly improve training stability and final loss.', 'type': 'bool'}, - 'schedulefree_c': {'title': 'Schedule free averaging strength', 'tooltip': 'Larger values = more responsive (shorter averaging window); smaller values = smoother (longer window). Set to 0 to disable and use the original Schedule-Free rule. Short small batches (≈6-12); long/large-batch (≈50-200).', 'type': 'float'}, - 'ns_steps': {'title': 'Newton-Schulz Iterations', 'tooltip': 'Controls the number of iterations for update orthogonalization. Higher values improve the updates quality but make each step slower. Lower values are faster per step but may be less effective.', 'type': 'int'}, - 'MuonWithAuxAdam': {'title': 'MuonWithAuxAdam', 'tooltip': 'Whether to use the standard way of Muon. Non-hidden layers fallback to ADAMW, and MUON takes the rest. Note: The auxiliary Adam (ADAMW) is typically only relevant for training "full" LoRA (LoRA for all layers) or full finetune and is irrelevant for most common LoRA use cases.', 'type': 'bool'}, - 'muon_hidden_layers': {'title': 'Hidden Layers', 'tooltip': 'Comma-separated list of hidden layers to train using Muon. Regular expressions (if toggled) are supported. Any model layer with a matching name will be trained using Muon. If None is provided it will default to using automatic way of finding hidden layers.', 'type': 'str'}, - 'muon_adam_regex': {'title': 'Use Regex', 'tooltip': 'Whether to use regular expressions for hidden layers.', 'type': 'bool'}, - 'muon_adam_lr': {'title': 'Auxiliary Adam LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer. If empty, it will use the main learning rate.', 'type': 'float'}, - 'muon_te1_adam_lr': {'title': 'AuxAdam TE1 LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer for the first text encoder. If empty, it will use the Auxiliary Adam LR.', 'type': 'float'}, - 'muon_te2_adam_lr': {'title': 'AuxAdam TE2 LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer for the second text encoder. If empty, it will use the Auxiliary Adam LR.', 'type': 'float'}, - 'rms_rescaling': {'title': 'RMS Rescaling', 'tooltip': 'Muon already scales its updates to approximate and use the same learning rate (LR) as Adam. This option integrates a more accurate method to match the Adam LR, but it is slower.', 'type': 'bool'}, - 'normuon_variant': {'title': 'NorMuon Variant', 'tooltip': 'Enables the NorMuon optimizer variant, which combines Muon orthogonalization with per-neuron adaptive learning rates for better convergence and balanced parameter updates. Costs only one scalar state buffer per parameter group, size few KBs, maintaining high memory efficiency.', 'type': 'bool'}, - 'beta2_normuon': {'title': 'NorMuon Beta2', 'tooltip': 'Exponential decay rate for the neuron-wise second-moment estimator in NorMuon (analogous to Adams beta2). Controls how past squared updates influence current normalization.', 'type': 'float'}, - 'low_rank_ortho': {'title': 'Low-rank Orthogonalization', 'tooltip': 'Use low-rank orthogonalization to accelerate Muon by orthogonalizing only in a low-dimensional subspace, improving speed and noise robustness.', 'type': 'bool'}, - 'ortho_rank': {'title': 'Ortho Rank', 'tooltip': 'Target rank for low-rank orthogonalization. Controls the dimensionality of the subspace used for efficient and noise-robust orthogonalization.', 'type': 'int'}, - 'accelerated_ns': {'title': 'Accelerated Newton-Schulz', 'tooltip': 'Applies an enhanced Newton-Schulz variant that replaces heuristic coefficients with optimal coefficients derived at each step. This improves performance and convergence by reducing the number of required operations.', 'type': 'bool'}, - 'cautious_wd': {'title': 'Cautious Weight Decay', 'tooltip': 'Applies weight decay only to parameter coordinates whose signs align with the optimizer update direction. This preserves the original optimization objective while still benefiting from regularization effects, leading to improved convergence and better final performance.', 'type': 'bool'}, - 'approx_mars': {'title': 'Approx MARS-M', 'tooltip': 'Enables Approximated MARS-M, a variance reduction technique. It uses the previous step\'s gradient to correct the current update, leading to lower losses and improved convergence stability. This requires additional state to store the previous gradient.', 'type': 'bool'}, - 'auto_kappa_p': {'title': 'Auto Lion-K', 'tooltip': 'Automatically determines the optimal P-value based on layer dimensions. Uses p=2.0 (Spherical) for 4D (Conv) tensors for stability and rotational invariance, and p=1.0 (Sign) for 2D (Linear) tensors for sparsity. Overrides the manual P-value. Recommend for unet models.', 'type': 'bool'}, - 'compile': {'title': 'Compiled Optimizer', 'tooltip': 'Enables PyTorch compilation for the optimizer internal step logic. This is intended to improve performance by allowing PyTorch to fuse operations and optimize the computational graph.', 'type': 'bool'}, - } - # @formatter:on - - if not self.winfo_exists(): # check if this window isn't open - return - - selected_optimizer = self.train_config.optimizer.optimizer - - # Extract the keys for the selected optimizer - for index, key in enumerate(OPTIMIZER_DEFAULT_PARAMETERS[selected_optimizer].keys()): - if key not in KEY_DETAIL_MAP: - continue - arg_info = KEY_DETAIL_MAP[key] - - title = arg_info['title'] - tooltip = arg_info['tooltip'] - type = arg_info['type'] - - row = (index // 2) + 1 - col = 3 * (index % 2) - - components.label(master, row, col, title, tooltip=tooltip) - - if key == 'MuonWithAuxAdam': - frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=row, column=col + 1, columnspan=2, sticky="ew", padx=0, pady=0) - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - - components.switch(frame, 0, 0, self.optimizer_ui_state, key, command=self.update_user_pref) - - self.muon_adam_button = components.button( - frame, 0, 1, "...", self.open_muon_adam_window, - tooltip="Configure the auxiliary AdamW_adv optimizer", - width=20, padx=5 ) - self.toggle_muon_adam_button() - elif type != 'bool': - components.entry(master, row, col + 1, self.optimizer_ui_state, key, - command=self.update_user_pref) - else: - components.switch(master, row, col + 1, self.optimizer_ui_state, key, - command=self.update_user_pref) - - def update_user_pref(self, *args): - update_optimizer_config(self.train_config) - self.toggle_muon_adam_button() - - def on_optimizer_change(self, *args): - optimizer_config = change_optimizer(self.train_config) - self.ui_state.get_var("optimizer").update(optimizer_config) - - self.clear_dynamic_ui(self.frame) - self.create_dynamic_ui(self.frame) - - def load_defaults(self, *args): - optimizer_config = load_optimizer_defaults(self.train_config) - self.ui_state.get_var("optimizer").update(optimizer_config) - - def on_window_close(self): - self.destroy() - - def toggle_muon_adam_button(self): - if self.muon_adam_button and self.muon_adam_button.winfo_exists(): - muon_with_adam = self.optimizer_ui_state.get_var("MuonWithAuxAdam").get() - self.muon_adam_button.configure(state="normal" if muon_with_adam else "disabled") - - def open_muon_adam_window(self): - current_optimizer = self.train_config.optimizer.optimizer + def on_close(self): + update_optimizer_config(self.config) + def prepare_muon_adam_config(self) -> tuple['TrainOptimizerConfig', Optimizer]: + current_optimizer = self.config.optimizer.optimizer adam_config = TrainOptimizerConfig.default_values() - current_state = self.train_config.optimizer.muon_adam_config + current_state = self.config.optimizer.muon_adam_config if current_optimizer == Optimizer.MUON: defaults = MUON_AUX_ADAM_DEFAULTS @@ -281,8 +45,7 @@ def open_muon_adam_window(self): # Should not happen if TrainConfig defines it as dict, but for safety adam_config = current_state - temp_adam_ui_state = UIState(self, adam_config) - window = MuonAdamWindow(self, self.train_config, temp_adam_ui_state, current_optimizer) - self.wait_window(window) + return adam_config, current_optimizer - self.train_config.optimizer.muon_adam_config = adam_config.to_dict() + def save_muon_adam_config(self, adam_config: 'TrainOptimizerConfig'): + self.config.optimizer.muon_adam_config = adam_config.to_dict() diff --git a/modules/ui/ProfilingWindowController.py b/modules/ui/ProfilingWindowController.py new file mode 100644 index 000000000..4e1f2d980 --- /dev/null +++ b/modules/ui/ProfilingWindowController.py @@ -0,0 +1,25 @@ +import faulthandler + + +class ProfilingWindowController: + def __init__(self): + self.view = None + + def create_window(self, parent, view_cls): + self.view = view_cls(parent, self) + return self.view + + def dump_stack(self): + with open('stacks.txt', 'w') as f: + faulthandler.dump_traceback(f) + self.view.set_message('Stack dumped to stacks.txt') + + def start_profiler(self): + from scalene import scalene_profiler + scalene_profiler.start() + self.view.set_profiling_active(True) + + def end_profiler(self): + from scalene import scalene_profiler + scalene_profiler.stop() + self.view.set_profiling_active(False) diff --git a/modules/ui/SampleFrameController.py b/modules/ui/SampleFrameController.py new file mode 100644 index 000000000..474c52ab8 --- /dev/null +++ b/modules/ui/SampleFrameController.py @@ -0,0 +1,17 @@ +from modules.util.config.SampleConfig import SampleConfig +from modules.util.enum.ModelType import ModelType + + +class SampleFrameController: + def __init__(self, sample: SampleConfig, model_type: ModelType): + self.sample = sample + self.model_type = model_type + + def is_flow_matching(self) -> bool: + return self.model_type.is_flow_matching() + + def is_inpainting_model(self) -> bool: + return self.model_type.has_conditioning_image_input() + + def is_video_model(self) -> bool: + return self.model_type.is_video_model() diff --git a/modules/ui/SampleParamsWindowController.py b/modules/ui/SampleParamsWindowController.py new file mode 100644 index 000000000..abe7c8b33 --- /dev/null +++ b/modules/ui/SampleParamsWindowController.py @@ -0,0 +1,8 @@ +from modules.util.config.SampleConfig import SampleConfig +from modules.util.enum.ModelType import ModelType + + +class SampleParamsWindowController: + def __init__(self, sample: SampleConfig, model_type: ModelType | None = None): + self.sample = sample + self.model_type = model_type diff --git a/modules/ui/SampleWindowController.py b/modules/ui/SampleWindowController.py index 0f91ad2fa..d1d24d643 100644 --- a/modules/ui/SampleWindowController.py +++ b/modules/ui/SampleWindowController.py @@ -1,49 +1,34 @@ -import contextlib import copy import os -import tkinter as tk -import traceback from modules.model.BaseModel import BaseModel from modules.modelSampler.BaseModelSampler import ( BaseModelSampler, - ModelSamplerOutput, ) -from modules.ui.SampleFrame import SampleFrame from modules.util import create from modules.util.callbacks.TrainCallbacks import TrainCallbacks from modules.util.commands.TrainCommands import TrainCommands from modules.util.config.SampleConfig import SampleConfig from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.EMAMode import EMAMode -from modules.util.enum.FileType import FileType from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.time_util import get_string_timestamp -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState import torch -import customtkinter as ctk -from PIL import Image - -class SampleWindow(ctk.CTkToplevel): +class SampleWindowController: def __init__( self, - parent, train_config: TrainConfig, use_external_model: bool, callbacks: TrainCallbacks | None = None, commands: TrainCommands | None = None, - *args, **kwargs ): - super().__init__(parent, *args, **kwargs) - - self.title("Sample") - self.geometry("1200x800") - self.resizable(True, True) + self.current_train_config = train_config + self.use_external_model = use_external_model + self.callbacks = callbacks + self.commands = commands if not use_external_model: self.initial_train_config = TrainConfig.default_values().from_dict(train_config.to_dict()) @@ -55,52 +40,19 @@ def __init__( #TODO why is there a current_train_config and an initial_train_config? #current_train_config doesn't seem to ever change - self.current_train_config = train_config - self.callbacks = callbacks - self.commands = commands # get model specific defaults model_type = train_config.model_type self.sample = SampleConfig.default_values(model_type) - self.ui_state = UIState(self, self.sample) - if use_external_model: - self.callbacks.set_on_sample_custom(self.__update_preview) - self.callbacks.set_on_update_sample_custom_progress(self.__update_progress) - else: + if not use_external_model: self.model = None self.model_sampler = None - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=1) - self.grid_rowconfigure(2, weight=0) - self.grid_rowconfigure(3, weight=0) - self.grid_columnconfigure(0, weight=0) - self.grid_columnconfigure(1, weight=1) - - prompt_frame = SampleFrame(self, self.sample, self.ui_state, include_settings=False, model_type=model_type) - prompt_frame.grid(row=0, column=0, columnspan=2, padx=0, pady=0, sticky="nsew") + def get_model_type(self): + return self.current_train_config.model_type - settings_frame = SampleFrame(self, self.sample, self.ui_state, include_prompt=False, model_type=model_type) - settings_frame.grid(row=1, column=0, padx=0, pady=0, sticky="nsew") - - # image - self.image = ctk.CTkImage( - light_image=self.__dummy_image(), - size=(512, 512) - ) - - image_label = ctk.CTkLabel(master=self, text="", image=self.image, height=512, width=512) - image_label.grid(row=1, column=1, rowspan=3, sticky="nsew") - - self.progress = components.progress(self, 2, 0) - components.button(self, 3, 0, "sample", self.__sample) - - self.wait_visibility() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - def __load_model(self) -> BaseModel: + def load_model(self) -> BaseModel: model_loader = create.create_model_loader( model_type=self.initial_train_config.model_type, training_method=self.initial_train_config.training_method, @@ -148,7 +100,7 @@ def __load_model(self) -> BaseModel: return model - def __create_sampler(self, model: BaseModel) -> BaseModelSampler: + def create_sampler(self, model: BaseModel) -> BaseModelSampler: return create.create_model_sampler( train_device=torch.device(self.initial_train_config.train_device), temp_device=torch.device(self.initial_train_config.temp_device), @@ -157,22 +109,7 @@ def __create_sampler(self, model: BaseModel) -> BaseModelSampler: training_method=self.initial_train_config.training_method, ) - def __update_preview(self, sampler_output: ModelSamplerOutput): - if sampler_output.file_type == FileType.IMAGE: - image = sampler_output.data - self.image.configure( - light_image=image, - size=(image.width, image.height), - ) - - def __update_progress(self, progress: int, max_progress: int): - self.progress.set(progress / max_progress) - self.update() - - def __dummy_image(self) -> Image: - return Image.new(mode="RGB", size=(512, 512), color=(0, 0, 0)) - - def __sample(self): + def do_sample(self, on_sample, on_update_progress): sample = copy.copy(self.sample) if self.commands: @@ -180,8 +117,8 @@ def __sample(self): else: if self.model is None: # lazy initialization - self.model = self.__load_model() - self.model_sampler = self.__create_sampler(self.model) + self.model = self.load_model() + self.model_sampler = self.create_sampler(self.model) sample.from_train_config(self.current_train_config) @@ -205,23 +142,6 @@ def __sample(self): image_format=self.current_train_config.sample_image_format, video_format=self.current_train_config.sample_video_format, audio_format=self.current_train_config.sample_audio_format, - on_sample=self.__update_preview, - on_update_progress=self.__update_progress, + on_sample=on_sample, + on_update_progress=on_update_progress, ) - - def destroy(self): - try: - if hasattr(self, "_icon_image_ref"): - del self._icon_image_ref - - # Remove any pending after callbacks - for after_id in self.tk.call('after', 'info'): - with contextlib.suppress(tk.TclError, RuntimeError): - self.after_cancel(after_id) - - super().destroy() - except (tk.TclError, RuntimeError) as e: - print(f"Error destroying window: {e}") - except Exception as e: - print(f"Unexpected error destroying window: {e}") - traceback.print_exc() diff --git a/modules/ui/SamplingTabController.py b/modules/ui/SamplingTabController.py new file mode 100644 index 000000000..ef95c77fa --- /dev/null +++ b/modules/ui/SamplingTabController.py @@ -0,0 +1,14 @@ +from modules.ui.SampleParamsWindowController import SampleParamsWindowController +from modules.util.config.SampleConfig import SampleConfig +from modules.util.config.TrainConfig import TrainConfig + + +class SamplingTabController: + def __init__(self, config: TrainConfig): + self.train_config = config + + def create_new_element(self) -> SampleConfig: + return SampleConfig.default_values(self.train_config.model_type) + + def open_element_window(self, parent, sample_config, ui_state, view_cls): + return view_cls(parent, SampleParamsWindowController(sample_config, model_type=self.train_config.model_type), ui_state) diff --git a/modules/ui/SchedulerParamsWindowController.py b/modules/ui/SchedulerParamsWindowController.py new file mode 100644 index 000000000..363391b04 --- /dev/null +++ b/modules/ui/SchedulerParamsWindowController.py @@ -0,0 +1,17 @@ +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.LearningRateScheduler import LearningRateScheduler + + +class SchedulerParamsWindowController: + def __init__(self, config: TrainConfig): + self.config = config + + def is_custom_scheduler(self) -> bool: + return self.config.learning_rate_scheduler is LearningRateScheduler.CUSTOM + +class KvParamsController: + def __init__(self, train_config: TrainConfig): + self.train_config = train_config + + def create_new_element(self) -> dict[str, str]: + return {"key": "", "value": ""} diff --git a/modules/ui/TimestepDistributionWindowController.py b/modules/ui/TimestepDistributionWindowController.py index 21e41ce3e..682508599 100644 --- a/modules/ui/TimestepDistributionWindowController.py +++ b/modules/ui/TimestepDistributionWindowController.py @@ -1,21 +1,11 @@ -from modules.modelSetup.mixin.ModelSetupNoiseMixin import ( - ModelSetupNoiseMixin, -) +from modules.modelSetup.mixin.ModelSetupNoiseMixin import ModelSetupNoiseMixin from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.TimestepDistribution import TimestepDistribution -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState import torch from torch import Tensor -import customtkinter as ctk -from customtkinter import AppearanceModeTracker, ThemeManager -from matplotlib import pyplot as plt -from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg - class TimestepGenerator(ModelSetupNoiseMixin): @@ -59,128 +49,20 @@ def generate(self) -> Tensor: ) -class TimestepDistributionWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - config: TrainConfig, - ui_state: UIState, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) - - self.title("Timestep Distribution") - self.geometry("900x600") - self.resizable(True, True) - - self.config = config - self.ui_state = ui_state - self.image_preview_file_index = 0 - self.ax = None - self.canvas = None - - self.grid_rowconfigure(0, weight=1) - self.grid_columnconfigure(0, weight=1) - - frame = self.__content_frame(self) - frame.grid(row=0, column=0, sticky='nsew') - components.button(self, 1, 0, "ok", self.__ok) - - self.wait_visibility() - self.after(200, lambda: set_window_icon(self)) - self.grab_set() - self.focus_set() - - def __content_frame(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - frame.grid_rowconfigure(7, weight=1) - - # timestep distribution - components.label(frame, 0, 0, "Timestep Distribution", - tooltip="Selects the function to sample timesteps during training", - wide_tooltip=True) - components.options(frame, 0, 1, [str(x) for x in list(TimestepDistribution)], self.ui_state, - "timestep_distribution") - - # min noising strength - components.label(frame, 1, 0, "Min Noising Strength", - tooltip="Specifies the minimum noising strength used during training. This can help to improve composition, but prevents finer details from being trained") - components.entry(frame, 1, 1, self.ui_state, "min_noising_strength") - - # max noising strength - components.label(frame, 2, 0, "Max Noising Strength", - tooltip="Specifies the maximum noising strength used during training. This can be useful to reduce overfitting, but also reduces the impact of training samples on the overall image composition") - components.entry(frame, 2, 1, self.ui_state, "max_noising_strength") +class TimestepDistributionWindowController: + def __init__(self, config: TrainConfig): + self.train_config = config - # noising weight - components.label(frame, 3, 0, "Noising Weight", - tooltip="Controls the weight parameter of the timestep distribution function. Use the preview to see more details.") - components.entry(frame, 3, 1, self.ui_state, "noising_weight") + def get_distribution_options(self) -> list[str]: + return [str(x) for x in list(TimestepDistribution)] - # noising bias - components.label(frame, 4, 0, "Noising Bias", - tooltip="Controls the bias parameter of the timestep distribution function. Use the preview to see more details.") - components.entry(frame, 4, 1, self.ui_state, "noising_bias") - - # timestep shift - components.label(frame, 5, 0, "Timestep Shift", - tooltip="Shift the timestep distribution. Use the preview to see more details.") - components.entry(frame, 5, 1, self.ui_state, "timestep_shift") - - # dynamic timestep shifting - components.label(frame, 6, 0, "Dynamic Timestep Shifting", - tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Dynamic Timestep Shifting is not shown in the preview. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) - components.switch(frame, 6, 1, self.ui_state, "dynamic_timestep_shifting") - - - # plot - appearance_mode = AppearanceModeTracker.get_mode() - background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) - text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) - background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" - text_color = f"#{int(text_color[0]/256):x}{int(text_color[1]/256):x}{int(text_color[2]/256):x}" - - fig, ax = plt.subplots() - self.ax = ax - self.canvas = FigureCanvasTkAgg(fig, master=frame) - self.canvas.get_tk_widget().grid(row=0, column=3, rowspan=8) - - fig.set_facecolor(background_color) - ax.set_facecolor(background_color) - ax.spines['bottom'].set_color(text_color) - ax.spines['left'].set_color(text_color) - ax.spines['top'].set_color(text_color) - ax.spines['right'].set_color(text_color) - ax.tick_params(axis='x', colors=text_color, which="both") - ax.tick_params(axis='y', colors=text_color, which="both") - ax.xaxis.label.set_color(text_color) - ax.yaxis.label.set_color(text_color) - - self.__update_preview() - - # update button - components.button(frame, 8, 3, "Update Preview", command=self.__update_preview) - - frame.pack(fill="both", expand=1) - return frame - - def __update_preview(self): + def generate_preview_data(self) -> Tensor: generator = TimestepGenerator( - timestep_distribution=self.config.timestep_distribution, - min_noising_strength=self.config.min_noising_strength, - max_noising_strength=self.config.max_noising_strength, - noising_weight=self.config.noising_weight, - noising_bias=self.config.noising_bias, - timestep_shift=self.config.timestep_shift, + timestep_distribution=self.train_config.timestep_distribution, + min_noising_strength=self.train_config.min_noising_strength, + max_noising_strength=self.train_config.max_noising_strength, + noising_weight=self.train_config.noising_weight, + noising_bias=self.train_config.noising_bias, + timestep_shift=self.train_config.timestep_shift, ) - - self.ax.cla() - self.ax.hist(generator.generate(), bins=1000, range=(0, 999)) - self.canvas.draw() - - def __ok(self): - self.destroy() + return generator.generate() diff --git a/modules/ui/TopBarController.py b/modules/ui/TopBarController.py index 820fdb71a..31d3fc1be 100644 --- a/modules/ui/TopBarController.py +++ b/modules/ui/TopBarController.py @@ -1,260 +1,103 @@ -import json import os -import traceback import webbrowser -from collections.abc import Callable -from contextlib import suppress from modules.util import path_util -from modules.util.config.SecretsConfig import SecretsConfig from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.ModelType import ModelType from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.optimizer_util import change_optimizer from modules.util.path_util import write_json_atomic -from modules.util.ui import components, dialogs -from modules.util.ui.UIState import UIState -import customtkinter as ctk - -class TopBar: - def __init__( - self, - master, - train_config: TrainConfig, - ui_state: UIState, - change_model_type_callback: Callable[[ModelType], None], - change_training_method_callback: Callable[[TrainingMethod], None], - load_preset_callback: Callable[[], None], - ): - self.master = master - self.train_config = train_config - self.ui_state = ui_state - self.change_model_type_callback = change_model_type_callback - self.change_training_method_callback = change_training_method_callback - self.load_preset_callback = load_preset_callback - - self.dir = "training_presets" - - self.config_ui_data = { - "config_name": path_util.canonical_join(self.dir, "#.json") - } - 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.current_config = [] - - self.frame = ctk.CTkFrame(master=master, corner_radius=0) - self.frame.grid(row=0, column=0, sticky="nsew") - - self.training_method = None - - # title - components.app_title(self.frame, 0, 0) - - # dropdown - self.configs_dropdown = None - self.__create_configs_dropdown() - - # remove button - # TODO - # components.icon_button(self.frame, 0, 2, "-", self.__remove_config) - - # Wiki button - components.button(self.frame, 0, 4, "Wiki", self.open_wiki, width=50) - - # save button - components.button(self.frame, 0, 3, "Save config", self.__save_config, - tooltip="Save the current configuration in a custom preset", width=90) - - # padding - self.frame.grid_columnconfigure(5, weight=1) - - # model type - components.options_kv( - master=self.frame, - row=0, - column=6, - values=[ #TODO simplify - ("SD1.5", ModelType.STABLE_DIFFUSION_15), - ("SD1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), - ("SD2.0", ModelType.STABLE_DIFFUSION_20), - ("SD2.0 Inpainting", ModelType.STABLE_DIFFUSION_20_INPAINTING), - ("SD2.1", ModelType.STABLE_DIFFUSION_21), - ("SD3", ModelType.STABLE_DIFFUSION_3), - ("SD3.5", ModelType.STABLE_DIFFUSION_35), - ("SDXL", ModelType.STABLE_DIFFUSION_XL_10_BASE), - ("SDXL Inpainting", ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING), - ("Wuerstchen v2", ModelType.WUERSTCHEN_2), - ("Stable Cascade", ModelType.STABLE_CASCADE_1), - ("PixArt Alpha", ModelType.PIXART_ALPHA), - ("PixArt Sigma", ModelType.PIXART_SIGMA), - ("Flux Dev.1", ModelType.FLUX_DEV_1), - ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), - ("Flux 2 [Dev, Klein]", ModelType.FLUX_2), - ("Sana", ModelType.SANA), - ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), - ("HiDream Full", ModelType.HI_DREAM_FULL), - ("Chroma1", ModelType.CHROMA_1), - ("QwenImage", ModelType.QWEN), - ("Z-Image", ModelType.Z_IMAGE), - ("Ernie Image", ModelType.ERNIE), - ], - ui_state=self.ui_state, - var_name="model_type", - command=self.__change_model_type, - ) - - def __create_training_method(self): - if self.training_method: - self.training_method.destroy() - - values = [] +class TopBarController: + def __init__(self, config: TrainConfig): + self.train_config = config + + def get_model_types(self) -> list[tuple[str, ModelType]]: + return [ #TODO simplify + ("SD1.5", ModelType.STABLE_DIFFUSION_15), + ("SD1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), + ("SD2.0", ModelType.STABLE_DIFFUSION_20), + ("SD2.0 Inpainting", ModelType.STABLE_DIFFUSION_20_INPAINTING), + ("SD2.1", ModelType.STABLE_DIFFUSION_21), + ("SD3", ModelType.STABLE_DIFFUSION_3), + ("SD3.5", ModelType.STABLE_DIFFUSION_35), + ("SDXL", ModelType.STABLE_DIFFUSION_XL_10_BASE), + ("SDXL Inpainting", ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING), + ("Wuerstchen v2", ModelType.WUERSTCHEN_2), + ("Stable Cascade", ModelType.STABLE_CASCADE_1), + ("PixArt Alpha", ModelType.PIXART_ALPHA), + ("PixArt Sigma", ModelType.PIXART_SIGMA), + ("Flux Dev.1", ModelType.FLUX_DEV_1), + ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), + ("Flux 2 [Dev, Klein]", ModelType.FLUX_2), + ("Sana", ModelType.SANA), + ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), + ("HiDream Full", ModelType.HI_DREAM_FULL), + ("Chroma1", ModelType.CHROMA_1), + ("QwenImage", ModelType.QWEN), + ("Z-Image", ModelType.Z_IMAGE), + ("Ernie Image", ModelType.ERNIE), + ] + + def get_training_methods(self, model_type: ModelType) -> list[tuple[str, TrainingMethod]]: #TODO simplify - if self.train_config.model_type.is_stable_diffusion(): - values = [ + if model_type.is_stable_diffusion(): + return [ ("Fine Tune", TrainingMethod.FINE_TUNE), ("LoRA", TrainingMethod.LORA), ("Embedding", TrainingMethod.EMBEDDING), ("Fine Tune VAE", TrainingMethod.FINE_TUNE_VAE), ] - elif self.train_config.model_type.is_stable_diffusion_3() \ - or self.train_config.model_type.is_stable_diffusion_xl() \ - or self.train_config.model_type.is_wuerstchen() \ - or self.train_config.model_type.is_pixart() \ - or self.train_config.model_type.is_flux_1() \ - or self.train_config.model_type.is_sana() \ - or self.train_config.model_type.is_hunyuan_video() \ - or self.train_config.model_type.is_hi_dream() \ - or self.train_config.model_type.is_chroma(): - values = [ + elif model_type.is_stable_diffusion_3() \ + or model_type.is_stable_diffusion_xl() \ + or model_type.is_wuerstchen() \ + or model_type.is_pixart() \ + or model_type.is_flux_1() \ + or model_type.is_sana() \ + or model_type.is_hunyuan_video() \ + or model_type.is_hi_dream() \ + or model_type.is_chroma(): + return [ ("Fine Tune", TrainingMethod.FINE_TUNE), ("LoRA", TrainingMethod.LORA), ("Embedding", TrainingMethod.EMBEDDING), ] - elif self.train_config.model_type.is_qwen() \ - or self.train_config.model_type.is_z_image() \ - or self.train_config.model_type.is_flux_2() \ - or self.train_config.model_type.is_ernie(): - values = [ + elif model_type.is_qwen() \ + or model_type.is_z_image() \ + or model_type.is_flux_2() \ + or model_type.is_ernie(): + return [ ("Fine Tune", TrainingMethod.FINE_TUNE), ("LoRA", TrainingMethod.LORA), ] + return [] - # training method - self.training_method = components.options_kv( - master=self.frame, - row=0, - column=7, - values=values, - ui_state=self.ui_state, - var_name="training_method", - command=self.change_training_method_callback, - ) - - def __change_model_type(self, model_type: ModelType): - self.change_model_type_callback(model_type) - self.__create_training_method() - - def __create_configs_dropdown(self): - if self.configs_dropdown is not None: - self.configs_dropdown.grid_forget() - - self.configs_dropdown = components.options_kv( - self.frame, 0, 1, self.configs, self.config_ui_state, "config_name", self.__load_current_config - ) - - def __load_available_config_names(self): - if os.path.isdir(self.dir): - for path in os.listdir(self.dir): + def load_available_config_names(self, dir: str) -> list[tuple[str, str]]: + configs = [("", path_util.canonical_join(dir, "#.json"))] + if os.path.isdir(dir): + for path in os.listdir(dir): if path != "#.json": - path = path_util.canonical_join(self.dir, path) + path = path_util.canonical_join(dir, path) if path.endswith(".json") and os.path.isfile(path): name = os.path.basename(path) name = os.path.splitext(name)[0] - self.configs.append((name, path)) - self.configs.sort() + configs.append((name, path)) + configs.sort() + return configs - def __save_to_file(self, name) -> str: + def save_to_file(self, name) -> str: name = path_util.safe_filename(name) path = path_util.canonical_join("training_presets", f"{name}.json") - write_json_atomic(path, self.train_config.to_settings_dict(secrets=False)) - return path - def __save_secrets(self, path) -> str: + def save_secrets(self, path) -> str: write_json_atomic(path, self.train_config.secrets.to_dict()) return path def open_wiki(self): webbrowser.open("https://github.com/Nerogar/OneTrainer/wiki", new=0, autoraise=False) - def __save_new_config(self, name): - path = self.__save_to_file(name) - - is_new_config = name not in [x[0] for x in self.configs] - - if is_new_config: - self.configs.append((name, path)) - self.configs.sort() - - if self.config_ui_data["config_name"] != path_util.canonical_join(self.dir, f"{name}.json"): - self.config_ui_state.get_var("config_name").set(path_util.canonical_join(self.dir, f"{name}.json")) - - if is_new_config: - self.__create_configs_dropdown() - - def __save_config(self): - default_value = self.configs_dropdown.get() - while default_value.startswith('#'): - default_value = default_value[1:] - - dialogs.StringInputDialog( - parent=self.master, - title="name", - question="Config Name", - callback=self.__save_new_config, - default_value=default_value, - validate_callback=lambda x: not x.startswith("#") - ) - - def __load_current_config(self, filename): - try: - basename = os.path.basename(filename) - is_built_in_preset = basename.startswith("#") and basename != "#.json" - - with open(filename, "r") as f: - loaded_dict = json.load(f) - default_config = TrainConfig.default_values() - if is_built_in_preset: - # always assume built-in configs are saved in the most recent version - loaded_dict["__version"] = default_config.config_version - loaded_config = default_config.from_dict(loaded_dict).to_unpacked_config() - - with suppress(FileNotFoundError), open("secrets.json", "r") as f: - secrets_dict=json.load(f) - loaded_config.secrets = SecretsConfig.default_values().from_dict(secrets_dict) - - self.train_config.from_dict(loaded_config.to_dict()) - self.ui_state.update(loaded_config) - - optimizer_config = change_optimizer(self.train_config) - self.ui_state.get_var("optimizer").update(optimizer_config) - - self.load_preset_callback() - except FileNotFoundError: - pass - except Exception: - print(traceback.format_exc()) - - def __remove_config(self): - # TODO - pass - def save_default(self): - self.__save_to_file("#") - self.__save_secrets("secrets.json") + self.save_to_file("#") + self.save_secrets("secrets.json") diff --git a/modules/ui/TrainUIController.py b/modules/ui/TrainUIController.py index ba90d2e64..15d1eaff6 100644 --- a/modules/ui/TrainUIController.py +++ b/modules/ui/TrainUIController.py @@ -9,47 +9,25 @@ import time import traceback import webbrowser -from collections.abc import Callable from contextlib import suppress from pathlib import Path -from tkinter import filedialog, messagebox import scripts.generate_debug_report -from modules.ui.AdditionalEmbeddingsTab import AdditionalEmbeddingsTab -from modules.ui.CaptionUI import CaptionUI -from modules.ui.CloudTab import CloudTab -from modules.ui.ConceptTab import ConceptTab -from modules.ui.ConvertModelUI import ConvertModelUI -from modules.ui.LoraTab import LoraTab -from modules.ui.ModelTab import ModelTab -from modules.ui.ProfilingWindow import ProfilingWindow -from modules.ui.SampleWindow import SampleWindow -from modules.ui.SamplingTab import SamplingTab -from modules.ui.TopBar import TopBar -from modules.ui.TrainingTab import TrainingTab -from modules.ui.VideoToolUI import VideoToolUI +from modules.ui.BaseTrainUIView import BaseTrainUIView +from modules.ui.CaptionUIController import CaptionUIController +from modules.ui.ConvertModelUIController import ConvertModelUIController +from modules.ui.SampleWindowController import SampleWindowController +from modules.ui.VideoToolUIController import VideoToolUIController from modules.util import create from modules.util.callbacks.TrainCallbacks import TrainCallbacks from modules.util.commands.TrainCommands import TrainCommands from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.DataType import DataType -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.PathIOType import PathIOType -from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState from modules.util.ui.validation import flush_and_validate_all import torch -import customtkinter as ctk -from customtkinter import AppearanceModeTracker - # 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": @@ -57,547 +35,25 @@ # 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): - set_step_progress: Callable[[int, int], None] - set_epoch_progress: Callable[[int, int], None] - - status_label: ctk.CTkLabel | None - training_button: ctk.CTkButton | None - training_callbacks: TrainCallbacks | None - training_commands: TrainCommands | None - - _TRAIN_BUTTON_STYLES = { - "idle": { - "text": "Start Training", - "state": "normal", - "fg_color": "#198754", - "hover_color": "#146c43", - "text_color": "white", - "text_color_disabled": "white", - }, - "running": { - "text": "Stop Training", - "state": "normal", - "fg_color": "#dc3545", - "hover_color": "#bb2d3b", - "text_color": "white", - }, - "stopping": { - "text": "Stopping...", - "state": "disabled", - "fg_color": "#dc3545", - "hover_color": "#dc3545", - "text_color": "white", - "text_color_disabled": "white", - }, - } - - def __init__(self): - super().__init__() - - self.title("OneTrainer") - self.geometry("1100x740") - - self.after(100, lambda: self._set_icon()) - - # more efficient version of ctk.set_appearance_mode("System"), which retrieves the system theme on each main loop iteration - ctk.set_appearance_mode("Light" if AppearanceModeTracker.detect_appearance_mode() == 0 else "Dark") - ctk.set_default_color_theme("blue") - - self.train_config = TrainConfig.default_values() - self.ui_state = UIState(self, self.train_config) - - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=1) - self.grid_rowconfigure(2, weight=0) - self.grid_columnconfigure(0, weight=1) - - self.status_label = None - self.eta_label = None - self.training_button = None - self.export_button = None - self.tabview = None - - self.model_tab = None - self.training_tab = None - self.lora_tab = None - self.cloud_tab = None - self.additional_embeddings_tab = None - - self.top_bar_component = self.top_bar(self) - self.content_frame(self) - self.bottom_bar(self) - self.training_thread = None - self.training_callbacks = None - self.training_commands = None +class TrainUIController: + def __init__(self, config: TrainConfig): + self.train_config = config + self.view: BaseTrainUIView | None = None + self.training_thread = None + self.training_callbacks: TrainCallbacks | None = None + self.training_commands: TrainCommands | None = None self.always_on_tensorboard_subprocess = None - self.current_workspace_dir = self.train_config.workspace_dir - self._check_start_always_on_tensorboard() - - self.workspace_dir_trace_id = self.ui_state.add_var_trace("workspace_dir", self._on_workspace_dir_change_trace) - - # Persistent profiling window. - self.profiling_window = ProfilingWindow(self) - - self.protocol("WM_DELETE_WINDOW", self.__close) - - def __close(self): - self.top_bar_component.save_default() - self._stop_always_on_tensorboard() - if hasattr(self, 'workspace_dir_trace_id'): - self.ui_state.remove_var_trace("workspace_dir", self.workspace_dir_trace_id) - self.quit() - - def top_bar(self, master): - return TopBar( - master, - self.train_config, - self.ui_state, - self.change_model_type, - self.change_training_method, - self.load_preset, - ) - - def _set_icon(self): - """Set the window icon safely after window is ready""" - set_window_icon(self) - - def bottom_bar(self, master): - frame = ctk.CTkFrame(master=master, corner_radius=0) - frame.grid(row=2, column=0, sticky="nsew") - - self.set_step_progress, self.set_epoch_progress = components.double_progress(frame, 0, 0, "step", "epoch") - - # status + ETA container - self.status_frame = ctk.CTkFrame(frame, corner_radius=0, fg_color="transparent") - self.status_frame.grid(row=0, column=1, sticky="w") - self.status_frame.grid_rowconfigure(0, weight=0) - self.status_frame.grid_rowconfigure(1, weight=0) - self.status_frame.grid_columnconfigure(0, weight=1) - - self.status_label = components.label(self.status_frame, 0, 0, "", pad=0, - tooltip="Current status of the training run") - self.eta_label = components.label(self.status_frame, 1, 0, "", pad=0) - - # padding - frame.grid_columnconfigure(2, weight=1) - - - # export button - self.export_button = components.button(frame, 0, 3, "Export", self.export_training, - width=60, padx=5, pady=(15, 0), - tooltip="Export the current configuration as a script to run without a UI") - - # debug button - components.button(frame, 0, 4, "Debug", self.generate_debug_package, - width=60, padx=(5, 25), pady=(15, 0), - tooltip="Generate a zip file with config.json, debug_report.log and settings diff, use this to report bugs or issues") - - # tensorboard button - components.button(frame, 0, 5, "Tensorboard", self.open_tensorboard, - width=100, padx=(0, 5), pady=(15, 0)) - - # training button - self.training_button = components.button(frame, 0, 6, "Start Training", self.start_training, - padx=(5, 20), pady=(15, 0)) - self._set_training_button_style("idle") # centralized styling - - return frame - - def content_frame(self, master): - frame = ctk.CTkFrame(master=master, corner_radius=0) - frame.grid(row=1, column=0, sticky="nsew") - - frame.grid_rowconfigure(0, weight=1) - frame.grid_columnconfigure(0, weight=1) - - self.tabview = ctk.CTkTabview(frame) - self.tabview.grid(row=0, column=0, sticky="nsew") - - 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) - - return frame - - def create_general_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # workspace dir - components.label(frame, 0, 0, "Workspace Directory", - tooltip="The directory where all files of this training run are saved") - components.path_entry(frame, 0, 1, self.ui_state, "workspace_dir", mode="dir", command=self._on_workspace_dir_change) - - # cache dir - components.label(frame, 0, 2, "Cache Directory", - tooltip="The directory where cached data is saved") - components.path_entry(frame, 0, 3, self.ui_state, "cache_dir", mode="dir") - - # continue from previous backup - components.label(frame, 2, 0, "Continue from last backup", - tooltip="Automatically continues training from the last backup saved in /backup") - components.switch(frame, 2, 1, self.ui_state, "continue_last_backup") - - # only cache - components.label(frame, 2, 2, "Only Cache", - tooltip="Only populate the cache, without any training") - components.switch(frame, 2, 3, self.ui_state, "only_cache") - - # TODO: In Phase 4 rework the general tab. - # prevent overwrites - components.label(frame, 3, 0, "Prevent Overwrites", - tooltip="When enabled, output paths that already exist on disk will be flagged as invalid to avoid accidental overwrites") - components.switch(frame, 3, 1, self.ui_state, "prevent_overwrites") - - # debug - components.label(frame, 4, 0, "Debug mode", - tooltip="Save debug information during the training into the debug directory") - components.switch(frame, 4, 1, self.ui_state, "debug_mode") - - components.label(frame, 4, 2, "Debug Directory", - tooltip="The directory where debug data is saved") - components.path_entry(frame, 4, 3, self.ui_state, "debug_dir", mode="dir", io_type=PathIOType.OUTPUT) - - # 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", - tooltip="Port to use for Tensorboard link") - components.entry(frame, 7, 3, self.ui_state, "tensorboard_port") - - - # validation - components.label(frame, 8, 0, "Validation", - tooltip="Enable validation steps and add new graph in tensorboard") - components.switch(frame, 8, 1, self.ui_state, "validation") - - components.label(frame, 8, 2, "Validate after", - tooltip="The interval used when validate training") - components.time_entry(frame, 8, 3, self.ui_state, "validate_after", "validate_after_unit") - - # device - components.label(frame, 10, 0, "Dataloader Threads", - tooltip="Number of threads used for the data loader. Increase if your GPU has room during caching, decrease if it's going out of memory during caching.") - components.entry(frame, 10, 1, self.ui_state, "dataloader_threads", required=True) - - components.label(frame, 11, 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", required=True) - - components.label(frame, 12, 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", - 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.label(frame, 13, 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, - "gradient_reduce_precision") - - components.label(frame, 13, 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.label(frame, 14, 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)", - 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.label(frame, 15, 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") - - frame.pack(fill="both", expand=1) - return frame - - def create_model_tab(self, master): - return ModelTab(master, self.train_config, self.ui_state) - - def create_data_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) - - # aspect ratio bucketing - components.label(frame, 0, 0, "Aspect Ratio Bucketing", - tooltip="Aspect ratio bucketing enables training on images with different aspect ratios") - components.switch(frame, 0, 1, self.ui_state, "aspect_ratio_bucketing") - - # latent caching - components.label(frame, 1, 0, "Latent Caching", - tooltip="Caching of intermediate training data that can be re-used between epochs") - components.switch(frame, 1, 1, self.ui_state, "latent_caching") - - # clear cache before training - components.label(frame, 2, 0, "Clear cache before training", - tooltip="Clears the cache directory before starting to train. Only disable this if you want to continue using the same cached data. Disabling this can lead to errors, if other settings are changed during a restart") - components.switch(frame, 2, 1, self.ui_state, "clear_cache_before_training") - - frame.pack(fill="both", expand=1) - return frame - - def create_concepts_tab(self, master): - return ConceptTab(master, self.train_config, self.ui_state) - - def create_training_tab(self, master) -> TrainingTab: - return TrainingTab(master, self.train_config, self.ui_state) - - def create_cloud_tab(self, master) -> CloudTab: - return CloudTab(master, self.train_config, self.ui_state,parent=self) - - def create_sampling_tab(self, master): - master.grid_rowconfigure(0, weight=0) - master.grid_rowconfigure(1, weight=1) - master.grid_columnconfigure(0, weight=1) - - # sample after - top_frame = ctk.CTkFrame(master=master, corner_radius=0) - top_frame.grid(row=0, column=0, sticky="nsew") - sub_frame = ctk.CTkFrame(master=top_frame, corner_radius=0, fg_color="transparent") - sub_frame.grid(row=1, column=0, sticky="nsew", columnspan=6) - - components.label(top_frame, 0, 0, "Sample After", - tooltip="The interval used when automatically sampling from the model during training") - components.time_entry(top_frame, 0, 1, self.ui_state, "sample_after", "sample_after_unit") - - components.label(top_frame, 0, 2, "Skip First", - tooltip="Start sampling automatically after this interval has elapsed.") - components.entry(top_frame, 0, 3, self.ui_state, "sample_skip_first", width=50, sticky="nw") - - components.label(top_frame, 0, 4, "Format", - tooltip="File Format used when saving samples") - components.options_kv(top_frame, 0, 5, [ - ("PNG", ImageFormat.PNG), - ("JPG", ImageFormat.JPG), - ], self.ui_state, "sample_image_format") - - components.button(top_frame, 0, 6, "sample now", self.sample_now) - - components.button(top_frame, 0, 7, "manual sample", self.open_manual_sample_window ) - - components.label(sub_frame, 0, 0, "Non-EMA Sampling", - tooltip="Whether to include non-ema sampling when using ema.") - components.switch(sub_frame, 0, 1, self.ui_state, "non_ema_sampling") - - components.label(sub_frame, 0, 2, "Samples to Tensorboard", - tooltip="Whether to include sample images in the Tensorboard output.") - components.switch(sub_frame, 0, 3, self.ui_state, "samples_to_tensorboard") - - # table - frame = ctk.CTkFrame(master=master, corner_radius=0) - frame.grid(row=1, column=0, sticky="nsew") - - return SamplingTab(frame, self.train_config, self.ui_state) - - def create_backup_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) - - # backup after - components.label(frame, 0, 0, "Backup After", - tooltip="The interval used when automatically creating model backups during training") - components.time_entry(frame, 0, 1, self.ui_state, "backup_after", "backup_after_unit") - - # backup now - components.button(frame, 0, 3, "backup now", self.backup_now) - - # rolling backup - components.label(frame, 1, 0, "Rolling Backup", - tooltip="If rolling backups are enabled, older backups are deleted automatically") - components.switch(frame, 1, 1, self.ui_state, "rolling_backup") - - # rolling backup count - components.label(frame, 1, 3, "Rolling Backup Count", - tooltip="Defines the number of backups to keep if rolling backups are enabled") - components.entry(frame, 1, 4, self.ui_state, "rolling_backup_count") - - # backup before save - components.label(frame, 2, 0, "Backup Before Save", - tooltip="Create a full backup before saving the final model") - components.switch(frame, 2, 1, self.ui_state, "backup_before_save") - - # save after - components.label(frame, 3, 0, "Save Every", - tooltip="The interval used when automatically saving the model during training") - components.time_entry(frame, 3, 1, self.ui_state, "save_every", "save_every_unit") - - # save now - components.button(frame, 3, 3, "save now", self.save_now) - - # skip save - components.label(frame, 4, 0, "Skip First", - tooltip="Start saving automatically after this interval has elapsed") - components.entry(frame, 4, 1, self.ui_state, "save_skip_first", width=50, sticky="nw") - - # save filename prefix - components.label(frame, 5, 0, "Save Filename Prefix", - tooltip="The prefix for filenames used when saving the model during training") - components.entry(frame, 5, 1, self.ui_state, "save_filename_prefix") - - frame.pack(fill="both", expand=1) - return frame - - def embedding_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) - - # 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.path_entry( - frame, 0, 1, self.ui_state, "embedding.model_name", - mode="file", path_modifier=components.json_path_modifier - ) + self.current_workspace_dir = config.workspace_dir + self.start_time: float = 0.0 - # token count - components.label(frame, 1, 0, "Token count", - tooltip="The token count used when creating a new embedding. Leave empty to auto detect from the initial embedding text.") - components.entry(frame, 1, 1, self.ui_state, "embedding.token_count") - - # initial embedding text - components.label(frame, 2, 0, "Initial embedding text", - tooltip="The initial embedding text used when creating a new embedding") - components.entry(frame, 2, 1, self.ui_state, "embedding.initial_embedding_text") - - # embedding weight dtype - components.label(frame, 3, 0, "Embedding Weight Data Type", - tooltip="The Embedding weight data type used for training. This can reduce memory consumption, but reduces precision") - components.options_kv(frame, 3, 1, [ - ("float32", DataType.FLOAT_32), - ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "embedding_weight_dtype") - - # placeholder - components.label(frame, 4, 0, "Placeholder", - tooltip="The placeholder used when using the embedding in a prompt") - components.entry(frame, 4, 1, self.ui_state, "embedding.placeholder") - - # output embedding - components.label(frame, 5, 0, "Output embedding", - tooltip="Output embeddings are calculated at the output of the text encoder, not the input. This can improve results for larger text encoders and lower VRAM usage.") - components.switch(frame, 5, 1, self.ui_state, "embedding.is_output_embedding") - - frame.pack(fill="both", expand=1) - return frame - - def create_additional_embeddings_tab(self, master): - return AdditionalEmbeddingsTab(master, self.train_config, self.ui_state) - - def create_tools_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) - - # dataset - components.label(frame, 0, 0, "Dataset Tools", - tooltip="Open the captioning tool") - components.button(frame, 0, 1, "Open", self.open_dataset_tool) - - # video tools - components.label(frame, 1, 0, "Video Tools", - tooltip="Open the video tools") - components.button(frame, 1, 1, "Open", self.open_video_tool) - - # convert model - components.label(frame, 2, 0, "Convert Model Tools", - tooltip="Open the model conversion tool") - components.button(frame, 2, 1, "Open", self.open_convert_model_tool) - - # sample - components.label(frame, 3, 0, "Sampling Tool", - tooltip="Open the model sampling tool") - components.button(frame, 3, 1, "Open", self.open_sampling_tool) - - components.label(frame, 4, 0, "Profiling Tool", - tooltip="Open the profiling tools.") - components.button(frame, 4, 1, "Open", self.open_profiling_tool) - - frame.pack(fill="both", expand=1) - return frame - - def change_model_type(self, model_type: ModelType): - if self.model_tab: - self.model_tab.refresh_ui() - - if self.training_tab: - self.training_tab.refresh_ui() - - if self.lora_tab: - self.lora_tab.refresh_ui() - - def change_training_method(self, training_method: TrainingMethod): - if not self.tabview: - return - - if self.model_tab: - self.model_tab.refresh_ui() - - if training_method != TrainingMethod.LORA and "LoRA" in self.tabview._tab_dict: - self.tabview.delete("LoRA") - self.lora_tab = None - if training_method != TrainingMethod.EMBEDDING and "embedding" in self.tabview._tab_dict: - self.tabview.delete("embedding") - - if training_method == TrainingMethod.LORA and "LoRA" not in self.tabview._tab_dict: - self.lora_tab = LoraTab(self.tabview.add("LoRA"), self.train_config, self.ui_state) - if training_method == TrainingMethod.EMBEDDING and "embedding" not in self.tabview._tab_dict: - self.embedding_tab(self.tabview.add("embedding")) - - def load_preset(self): - if not self.tabview: - return - - if self.additional_embeddings_tab: - self.additional_embeddings_tab.refresh_ui() + def on_update_train_progress(self, train_progress: TrainProgress, max_step: int, max_epoch: int): + eta_str = self._calculate_eta_string(train_progress, max_step, max_epoch) + self.view.on_update_progress(train_progress.epoch_step, max_step, train_progress.epoch, max_epoch, eta_str) - def open_tensorboard(self): - webbrowser.open("http://localhost:" + str(self.train_config.tensorboard_port), new=0, autoraise=False) + def on_update_status(self, status: str): + self.view.on_update_status(status) def _calculate_eta_string(self, train_progress: TrainProgress, max_step: int, max_epoch: int) -> str | None: spent_total = time.monotonic() - self.start_time @@ -621,85 +77,136 @@ def _calculate_eta_string(self, train_progress: TrainProgress, max_step: int, ma else: return f"{seconds}s" - def set_eta_label(self, train_progress: TrainProgress, max_step: int, max_epoch: int): - eta_str = self._calculate_eta_string(train_progress, max_step, max_epoch) - if eta_str is not None: - self.eta_label.configure(text=f"ETA: {eta_str}") - else: - self.eta_label.configure(text="") + 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 delete_eta_label(self): - self.eta_label.configure(text="") + def _start_always_on_tensorboard(self): + if self.always_on_tensorboard_subprocess: + self._stop_always_on_tensorboard() - def on_update_train_progress(self, train_progress: TrainProgress, max_step: int, max_epoch: int): - self.set_step_progress(train_progress.epoch_step, max_step) - self.set_epoch_progress(train_progress.epoch, max_epoch) - self.set_eta_label(train_progress, max_step, max_epoch) + 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_update_status(self, status: str): - self.status_label.configure(text=status) + os.makedirs(Path(tensorboard_log_dir).absolute(), exist_ok=True) - def open_dataset_tool(self): - window = CaptionUI(self, None, False) - self.wait_window(window) + tensorboard_args = [ + tensorboard_executable, + "--logdir", + tensorboard_log_dir, + "--port", + str(self.train_config.tensorboard_port), + "--samples_per_plugin=images=100,scalars=10000", + ] - def open_video_tool(self): - window = VideoToolUI(self) - self.wait_window(window) + if self.train_config.tensorboard_expose: + tensorboard_args.append("--bind_all") - def open_convert_model_tool(self): - window = ConvertModelUI(self) - self.wait_window(window) + try: + self.always_on_tensorboard_subprocess = subprocess.Popen(tensorboard_args) + except Exception: + self.always_on_tensorboard_subprocess = None - def open_sampling_tool(self): - if not self.training_callbacks and not self.training_commands: - window = SampleWindow( - self, - use_external_model=False, - train_config=self.train_config, - ) - self.wait_window(window) - torch_gc() + 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 open_profiling_tool(self): - self.profiling_window.deiconify() + 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 - def generate_debug_package(self): - zip_path = filedialog.askdirectory( - initialdir=".", - title="Select Directory to Save Debug Package" - ) + if self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: + self._start_always_on_tensorboard() - if not zip_path: - return + 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 - zip_path = Path(zip_path) / "OneTrainer_debug_report.zip" + if self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: + self._start_always_on_tensorboard() - self.on_update_status("Generating debug package...") + 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() - try: - config_json_string = json.dumps(self.train_config.to_pack_dict(secrets=False)) - scripts.generate_debug_report.create_debug_package(str(zip_path), config_json_string) - self.on_update_status(f"Debug package saved to {zip_path.name}") - except Exception as e: - traceback.print_exc() - self.on_update_status(f"Error generating debug package: {e}") + def open_tensorboard(self): + webbrowser.open("http://localhost:" + str(self.train_config.tensorboard_port), new=0, autoraise=False) + + def open_dataset_tool(self, parent, view_cls): + return CaptionUIController(None, False).create_window(parent, view_cls) + def open_video_tool(self, parent, view_cls): + return VideoToolUIController().create_window(parent, view_cls) + + def open_convert_model_tool(self, parent, view_cls): + return ConvertModelUIController().create_window(parent, view_cls) + + def open_sampling_tool(self, parent, view_cls): + if not self.training_callbacks and not self.training_commands: + controller = SampleWindowController( + self.train_config, + use_external_model=False, + ) + window = view_cls(parent, controller) + parent.show_window(window) + torch_gc() - def open_manual_sample_window (self): + def open_manual_sample_window(self, parent, view_cls): training_callbacks = self.training_callbacks training_commands = self.training_commands if training_callbacks and training_commands: - window = SampleWindow( - self, - train_config=self.train_config, + controller = SampleWindowController( + self.train_config, use_external_model=True, callbacks=training_callbacks, commands=training_commands, ) - self.wait_window(window) - training_callbacks.set_on_sample_custom() + window = view_cls(parent, controller) + parent.show_window(window) + parent.connect_window_closed(window, lambda: training_callbacks.set_on_sample_custom()) + + def sample_now(self): + train_commands = self.training_commands + if train_commands: + train_commands.sample_default() + + def backup_now(self): + train_commands = self.training_commands + if train_commands: + train_commands.backup() + + def save_now(self): + train_commands = self.training_commands + if train_commands: + train_commands.save() + + def export_training(self, file_path: str): + with open(file_path, "w") as f: + json.dump(self.train_config.to_pack_dict(secrets=False), f, indent=4) + + def generate_debug_package(self, zip_path: Path): + self.view.on_update_status("Generating debug package...") + try: + config_json_string = json.dumps(self.train_config.to_pack_dict(secrets=False)) + scripts.generate_debug_report.create_debug_package(str(zip_path), config_json_string) + self.view.on_update_status(f"Debug package saved to {zip_path.name}") + except Exception as e: + traceback.print_exc() + self.view.on_update_status(f"Error generating debug package: {e}") def __training_thread_function(self): error_caught = False @@ -709,17 +216,17 @@ def __training_thread_function(self): on_update_status=self.on_update_status, ) - trainer = create.create_trainer(self.train_config, self.training_callbacks, self.training_commands, reattach=self.cloud_tab.reattach) + trainer = create.create_trainer(self.train_config, self.training_callbacks, self.training_commands, reattach=self.view.get_cloud_reattach()) try: trainer.start() if self.train_config.cloud.enabled: - self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) + self.view.sync_cloud_secrets() self.start_time = time.monotonic() trainer.train() except Exception: if self.train_config.cloud.enabled: - self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) + self.view.sync_cloud_secrets() error_caught = True traceback.print_exc() @@ -737,30 +244,23 @@ def __training_thread_function(self): self.on_update_status("Error: check the console for details") else: self.on_update_status("Stopped") - self.delete_eta_label() - # queue UI update on Tk main thread; _set_training_button_idle applies shared styles, avoid potential race/crash - self.after(0, self._set_training_button_idle) + # queue UI update on Tk main thread; on_training_stopped applies shared styles, avoid potential race/crash + self.view.schedule_on_main_thread(lambda: self.view.on_training_stopped(error_caught)) if self.train_config.tensorboard_always_on and not self.always_on_tensorboard_subprocess: - self.after(0, self._start_always_on_tensorboard) + self.view.schedule_on_main_thread(self._start_always_on_tensorboard) def start_training(self): if self.training_thread is None: - self.save_default() + self.view.save_default() - # --- pre-training validation gate --- errors = flush_and_validate_all() - if errors: - bullet_list = "\n".join(f"• {e}" for e in errors) - messagebox.showerror( - "Cannot Start Training", - f"Please fix the following errors before training:\n\n{bullet_list}", - ) + self.view.show_validation_errors(errors) return - self._set_training_button_running() + self.view.on_training_started() if self.train_config.tensorboard and not self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: self._stop_always_on_tensorboard() @@ -771,119 +271,6 @@ def start_training(self): self.training_thread = threading.Thread(target=self.__training_thread_function) self.training_thread.start() else: - self._set_training_button_stopping() + self.view.on_training_stopping() self.on_update_status("Stopping ...") self.training_commands.stop() - - 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() - - def export_training(self): - file_path = filedialog.asksaveasfilename(filetypes=[ - ("All Files", "*.*"), - ("json", "*.json"), - ], initialdir=".", initialfile="config.json") - - if file_path: - with open(file_path, "w") as f: - json.dump(self.train_config.to_pack_dict(secrets=False), f, indent=4) - - def sample_now(self): - train_commands = self.training_commands - if train_commands: - train_commands.sample_default() - - def backup_now(self): - train_commands = self.training_commands - if train_commands: - train_commands.backup() - - def save_now(self): - train_commands = self.training_commands - 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") - - 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", - ] - - if self.train_config.tensorboard_expose: - tensorboard_args.append("--bind_all") - - 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 _set_training_button_style(self, mode: str): - if not self.training_button: - return - style = self._TRAIN_BUTTON_STYLES.get(mode) - if not style: - return - self.training_button.configure(**style) - - def _set_training_button_idle(self): - self._set_training_button_style("idle") - - def _set_training_button_running(self): - self._set_training_button_style("running") - - def _set_training_button_stopping(self): - self._set_training_button_style("stopping") diff --git a/modules/ui/TrainingTabController.py b/modules/ui/TrainingTabController.py new file mode 100644 index 000000000..858ca0701 --- /dev/null +++ b/modules/ui/TrainingTabController.py @@ -0,0 +1,39 @@ + +from modules.ui.OffloadingWindowController import OffloadingWindowController +from modules.ui.OptimizerParamsWindowController import OptimizerParamsWindowController +from modules.ui.SchedulerParamsWindowController import SchedulerParamsWindowController +from modules.ui.TimestepDistributionWindowController import TimestepDistributionWindowController +from modules.util import create +from modules.util.config.TrainConfig import TrainConfig +from modules.util.optimizer_util import change_optimizer + + +class TrainingTabController: + def __init__(self, config: TrainConfig): + self.config = config + + def restore_optimizer_config(self, ui_state): + optimizer_config = change_optimizer(self.config) + ui_state.get_var("optimizer").update(optimizer_config) + + def get_layer_presets(self) -> dict: + cls = create.get_model_setup_class(self.config.model_type, self.config.training_method) + return cls.LAYER_PRESETS if cls is not None else {"full": []} + + def is_flow_matching(self) -> bool: + return self.config.model_type.is_flow_matching() + + def is_custom_scheduler_value(self, value: str) -> bool: + return value == "CUSTOM" + + def open_optimizer_params_window(self, parent, ui_state, view_cls): + return view_cls(parent, OptimizerParamsWindowController(self.config), ui_state) + + def open_scheduler_params_window(self, parent, ui_state, view_cls): + return view_cls(parent, SchedulerParamsWindowController(self.config), ui_state) + + def open_timestep_distribution_window(self, parent, ui_state, view_cls): + return view_cls(parent, TimestepDistributionWindowController(self.config), ui_state) + + def open_offloading_window(self, parent, ui_state, view_cls): + return view_cls(parent, OffloadingWindowController(self.config), ui_state) diff --git a/modules/ui/VideoToolUIController.py b/modules/ui/VideoToolUIController.py index c3291e6ea..294f458f4 100644 --- a/modules/ui/VideoToolUIController.py +++ b/modules/ui/VideoToolUIController.py @@ -6,338 +6,54 @@ import shlex import subprocess import threading -import webbrowser from fractions import Fraction -from tkinter import filedialog -from modules.util.image_util import load_image from modules.util.path_util import SUPPORTED_VIDEO_EXTENSIONS -from modules.util.ui import components import av -import customtkinter as ctk import cv2 import scenedetect from PIL import Image -class VideoToolUI(ctk.CTkToplevel): - def __init__( - self, - parent, - *args, **kwargs, - ): - ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) - - self.title("Video Tools") - self.geometry("600x720") - self.resizable(True, True) - self.wait_visibility() - self.focus_set() - - self.grid_rowconfigure(0, weight=1) - self.grid_rowconfigure(1, weight=0) - self.grid_columnconfigure(0, weight=1) - - tabview = ctk.CTkTabview(self) - tabview.grid(row=0, column=0, sticky="nsew") - - self.clip_extract_tab = self.__clip_extract_tab(tabview.add("extract clips")) - self.image_extract_tab = self.__image_extract_tab(tabview.add("extract images")) - self.video_download_tab = self.__video_download_tab(tabview.add("download")) - self.status_bar(self) - - def status_bar(self, master): - frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=1, column=0) - frame.grid_columnconfigure(0, weight=0, minsize=160) - frame.grid_columnconfigure(1, weight=0, minsize=300) - frame.grid_columnconfigure(2, weight=1) - - #create preview image - preview_path = "resources/icons/icon.png" - preview = load_image(preview_path, 'RGB') - preview.thumbnail((150, 150)) - self.preview_image= ctk.CTkImage(light_image=preview, size=preview.size) - self.preview_image_label = ctk.CTkLabel( - master=frame, text="Preview image", image=self.preview_image, height=150, width=150, - compound="top") - self.preview_image_label.grid(row=0, column=0, sticky="nw", padx=5, pady=5) - - #displays progress and messages that also go to terminal - self.status_label = ctk.CTkTextbox(master=frame, width=400, height=160, wrap="word", border_width=2) - self.status_label.insert(index="1.0", text="Current status") - self.status_label.configure(state="disabled") - self.status_label.grid(row=0, column=1, sticky="ne", padx=5, pady=5) - - def __clip_extract_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=120) - frame.grid_columnconfigure(1, weight=0, minsize=200) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # single video - components.label(frame, 0, 0, "Single Video", - tooltip="Link to single video file to process.") - self.clip_single_entry = ctk.CTkEntry(frame, width=190) - self.clip_single_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) - self.clip_single_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_file(self.clip_single_entry, - [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))] - )) - self.clip_single_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 0, 2, "Extract Single", - command=lambda: self.__extract_clips_button(False)) - - # time range - components.label(frame, 1, 0, " Time Range", - tooltip="Time range to limit selection for single video, \ - format as hour:minute:second, minute:second, or seconds.") - self.clip_time_start_entry = ctk.CTkEntry(frame, width=100) - self.clip_time_start_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.clip_time_start_entry.insert(0, "00:00:00") - self.clip_time_end_entry = ctk.CTkEntry(frame, width=100) - self.clip_time_end_entry.grid(row=1, column=1, sticky="e", padx=5, pady=5) - self.clip_time_end_entry.insert(0, "99:99:99") - - # directory of videos - components.label(frame, 2, 0, "Directory", - tooltip="Path to directory with multiple videos to process, including in subdirectories.") - self.clip_list_entry = ctk.CTkEntry(frame, width=190) - self.clip_list_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - self.clip_list_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.clip_list_entry)) - self.clip_list_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 2, 2, "Extract Directory", - command=lambda: self.__extract_clips_button(True)) - - # output directory - components.label(frame, 3, 0, "Output", - tooltip="Path to folder where extracted clips will be saved.") - self.clip_output_entry = ctk.CTkEntry(frame, width=190) - self.clip_output_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) - self.clip_output_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.clip_output_entry)) - self.clip_output_button.grid(row=3, column=1, sticky="e", padx=5, pady=5) - - # output to subdirectories - self.output_subdir_clip = ctk.BooleanVar(self, False) - components.label(frame, 4, 0, "Output to\nSubdirectories", - tooltip="If enabled, files are saved to subfolders based on filename and input directory. \ - Otherwise will all be saved to the top level of the output directory.") - self.output_subdir_clip_entry = ctk.CTkSwitch(frame, variable=self.output_subdir_clip, text="") - self.output_subdir_clip_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) - - # split at cuts - self.split_at_cuts = ctk.BooleanVar(self, False) - components.label(frame, 5, 0, "Split at Cuts", - tooltip="If enabled, detect cuts in the input video and split at those points. \ - Otherwise will split at any point, and clips may contain cuts.") - self.split_cuts_entry = ctk.CTkSwitch(frame, variable=self.split_at_cuts, text="") - self.split_cuts_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) - - # maximum length - components.label(frame, 6, 0, "Max Length (s)", - tooltip="Maximum length in seconds for saved clips, larger clips will be broken into multiple small clips.") - self.clip_length_entry = ctk.CTkEntry(frame, width=220) - self.clip_length_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) - self.clip_length_entry.insert(0, "3") - - # Set FPS - components.label(frame, 7, 0, "Set FPS", - tooltip="FPS to convert output videos to, set to 0 to keep original rate.") - self.clip_fps_entry = ctk.CTkEntry(frame, width=220) - self.clip_fps_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) - self.clip_fps_entry.insert(0, "24.0") - - # Remove borders - self.clip_bordercrop = ctk.BooleanVar(self, False) - components.label(frame, 8, 0, "Remove Borders", - tooltip="Remove black borders from output clip") - self.clip_bordercrop_entry = ctk.CTkSwitch(frame, variable=self.clip_bordercrop, text="") - self.clip_bordercrop_entry.grid(row=8, column=1, sticky="w", padx=5, pady=5) - - # Crop Variation - components.label(frame, 9, 0, "Crop Variation", - tooltip="Output clips will be randomly cropped to +- the base aspect ratio, \ - somewhat biased towards making square videos. Set to 0 to use only base aspect.") - self.clip_crop_entry = ctk.CTkEntry(frame, width=220) - self.clip_crop_entry.grid(row=9, column=1, sticky="w", padx=5, pady=5) - self.clip_crop_entry.insert(0, "0.2") - - # object filter - currently unused, may implement in future - # components.label(frame, 9, 0, "Object Filter", - # tooltip="Detect general features using Haar-Cascade classifier, and choose how to deal with clips where it is detected") - # components.options(frame, 9, 1, ["NONE", "FACE", "EYE", "BODY"], self.video_ui_state, "filter_object") - # components.options(frame, 9, 2, ["INCLUDE", "EXCLUDE", "SUBFOLDER"], self.video_ui_state, "filter_behavior") - - frame.pack(fill="both", expand=1) - return frame - - def __image_extract_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=120) - frame.grid_columnconfigure(1, weight=0, minsize=200) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # single video - components.label(frame, 0, 0, "Single Video", - tooltip="Link to single video file to process.") - self.image_single_entry = ctk.CTkEntry(frame, width=190) - self.image_single_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) - self.image_single_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_file(self.image_single_entry, - [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))] - )) - self.image_single_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 0, 2, "Extract Single", - command=lambda: self.__extract_images_button(False)) - - # time range - components.label(frame, 1, 0, " Time Range", - tooltip="Time range to limit selection for single video, \ - format as hour:minute:second, minute:second, or seconds.") - self.image_time_start_entry = ctk.CTkEntry(frame, width=100) - self.image_time_start_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.image_time_start_entry.insert(0, "00:00:00") - self.image_time_end_entry = ctk.CTkEntry(frame, width=100) - self.image_time_end_entry.grid(row=1, column=1, sticky="e", padx=5, pady=5) - self.image_time_end_entry.insert(0, "99:99:99") - - # directory of videos - components.label(frame, 2, 0, "Directory", - tooltip="Path to directory with multiple videos to process, including in subdirectories.") - self.image_list_entry = ctk.CTkEntry(frame, width=190) - self.image_list_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - self.image_list_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.image_list_entry)) - self.image_list_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 2, 2, "Extract Directory", - command=lambda: self.__extract_images_button(True)) - - # output directory - components.label(frame, 3, 0, "Output", - tooltip="Path to folder where extracted images will be saved.") - self.image_output_entry = ctk.CTkEntry(frame, width=190) - self.image_output_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) - self.image_output_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.image_output_entry)) - self.image_output_button.grid(row=3, column=1, sticky="e", padx=5, pady=5) - - # output to subdirectories - self.output_subdir_img = ctk.BooleanVar(self, False) - components.label(frame, 4, 0, "Output to\nSubdirectories", - tooltip="If enabled, files are saved to subfolders based on filename and input directory. \ - Otherwise will all be saved to the top level of the output directory.") - self.output_subdir_img_entry = ctk.CTkSwitch(frame, variable=self.output_subdir_img, text="") - self.output_subdir_img_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) - - # image capture rate - components.label(frame, 5, 0, "Images/sec", - tooltip="Number of images to capture per second of video. \ - Images will be taken at semi-random frames around the specified frequency.") - self.capture_rate_entry = ctk.CTkEntry(frame, width=220) - self.capture_rate_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) - self.capture_rate_entry.insert(0, "0.5") - - # blur removal - components.label(frame, 6, 0, "Blur Removal", - tooltip="Threshold for removal of blurry images, relative to all others. \ - For example at 0.2, the blurriest 20%% of the final selected frames will not be saved.") - self.blur_threshold_entry = ctk.CTkEntry(frame, width=220) - self.blur_threshold_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) - self.blur_threshold_entry.insert(0, "0.2") - - # Remove borders - self.image_bordercrop = ctk.BooleanVar(self, False) - components.label(frame, 7, 0, "Remove Borders", - tooltip="Remove black borders from output image") - self.image_bordercrop_entry = ctk.CTkSwitch(frame, variable=self.image_bordercrop, text="") - self.image_bordercrop_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) - - # Crop Variation - components.label(frame, 8, 0, "Crop Variation", - tooltip="Output images will be randomly cropped to +- the base aspect ratio, \ - somewhat biased towards making square images. Set to 0 to use only base sapect.") - self.image_crop_entry = ctk.CTkEntry(frame, width=220) - self.image_crop_entry.grid(row=8, column=1, sticky="w", padx=5, pady=5) - self.image_crop_entry.insert(0, "0.2") - - # # object filter - currently unused, may implement in future - # components.label(frame, 5, 0, "Object Filter", - # tooltip="Detect general features using Haar-Cascade classifier, and choose how to deal with clips where it is detected") - # components.options(frame, 5, 1, ["NONE", "FACE", "EYE", "BODY"], self.video_ui_state, "filter_object") - # components.options(frame, 5, 2, ["INCLUDE", "EXCLUDE", "SUBFOLDER"], self.video_ui_state, "filter_behavior") - - frame.pack(fill="both", expand=1) - return frame - - def __video_download_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=120) - frame.grid_columnconfigure(1, weight=0, minsize=200) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # link - components.label(frame, 0, 0, "Single Link", - tooltip="Link to video/playlist to download. Uses yt-dlp, supports youtube, twitch, instagram, and many other sites.") - self.download_link_entry = ctk.CTkEntry(frame, width=220) - self.download_link_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) - components.button(frame, 0, 2, "Download Link", command=lambda: self.__download_button(False)) - - # link list - components.label(frame, 1, 0, "Link List", - tooltip="Path to txt file with list of links separated by newlines.") - self.download_list_entry = ctk.CTkEntry(frame, width=190) - self.download_list_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.download_list_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_file(self.download_list_entry, [("Text file", ".txt")])) - self.download_list_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 1, 2, "Download List", command=lambda: self.__download_button(True)) - - # output directory - components.label(frame, 2, 0, "Output", - tooltip="Path to folder where downloaded videos will be saved.") - self.download_output_entry = ctk.CTkEntry(frame, width=190) - self.download_output_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - self.download_output_button = ctk.CTkButton(frame, width=30, text="...", command=lambda: self.__browse_for_dir(self.download_output_entry)) - self.download_output_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) - - # additional args - 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.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 --format mp4") - components.button(frame, 3, 2, "yt-dlp info", - command=lambda: webbrowser.open("https://github.com/yt-dlp/yt-dlp?tab=readme-ov-file#usage-and-options", new=0, autoraise=False)) - - frame.pack(fill="both", expand=1) - return frame - - def __browse_for_dir(self, entry_box): - # get the path from the user - path = filedialog.askdirectory() - # set the path to the entry box - # delete entry box text - entry_box.focus_set() - entry_box.delete(0, ctk.END) - entry_box.insert(0, path) - self.focus_set() - - def __browse_for_file(self, entry_box, filetypes): - # get the path from the user - path = filedialog.askopenfilename(filetypes=filetypes) - # set the path to the entry box - # delete entry box text - entry_box.focus_set() - entry_box.delete(0, ctk.END) - entry_box.insert(0, path) - self.focus_set() +class VideoToolUIController: + def __init__(self): + self.view = None + self.args = { + "clip_single": "", + "clip_list": "", + "clip_time_start": "00:00:00", + "clip_time_end": "99:99:99", + "clip_output": "", + "output_subdir_clip": False, + "split_cuts": False, + "clip_length": "3", + "clip_fps": "24.0", + "clip_bordercrop": False, + "clip_crop": "0.2", + "image_single": "", + "image_list": "", + "image_time_start": "00:00:00", + "image_time_end": "99:99:99", + "image_output": "", + "output_subdir_img": False, + "capture_rate": "0.5", + "blur_threshold": "0.2", + "image_bordercrop": False, + "image_crop": "0.2", + "download_link": "", + "download_list": "", + "download_output": "", + "download_args": "--quiet --no-warnings --progress --format mp4", + } + + def create_window(self, parent, view_cls): + self.view = view_cls(parent, self) + return self.view + + def __update_status(self, status_text: str): + print(status_text) + self.view.update_status(status_text) def __get_vid_paths(self, batch_mode: bool, input_path_single: str, input_path_dir: str): input_videos = [] @@ -382,9 +98,7 @@ def __get_vid_paths(self, batch_mode: bool, input_path_single: str, input_path_d def __run_in_thread(self, target, *args): """Clear status box and run target function in a daemon thread.""" - self.status_label.configure(state="normal") - self.status_label.delete(index1="1.0", index2="end") - self.status_label.configure(state="disabled") + self.view.clear_status() t = threading.Thread(target=target, args=args) t.daemon = True t.start() @@ -463,22 +177,20 @@ def find_main_contour(self, frame): h1, w1, _ = frame.shape return x1, y1, w1, h1 - def __extract_clips_button(self, batch_mode: bool): + def extract_clips_button(self, batch_mode: bool): self.__run_in_thread(self.__extract_clips_multi, batch_mode) def __extract_clips_multi(self, batch_mode: bool): - if not pathlib.Path(self.clip_output_entry.get()).is_dir() or self.clip_output_entry.get() == "": + p = self.args + if not pathlib.Path(p['clip_output']).is_dir() or p['clip_output'] == "": self.__update_status("Invalid output directory!") return # validate numeric inputs try: - max_length = float(self.clip_length_entry.get()) - crop_variation = float(self.clip_crop_entry.get()) - target_fps = float(self.clip_fps_entry.get()) - input_single_entry = self.clip_single_entry.get() - input_multiple_entry = self.clip_list_entry.get() - output_entry = self.clip_output_entry.get() + max_length = float(p['clip_length']) + crop_variation = float(p['clip_crop']) + target_fps = float(p['clip_fps']) except ValueError: self.__update_status("Invalid numeric input for Max Length, Crop Variation, or FPS.") return @@ -492,26 +204,26 @@ def __extract_clips_multi(self, batch_mode: bool): self.__update_status("Crop Variation must be between 0.0 and 1.0.") return - input_videos = self.__get_vid_paths(batch_mode, input_single_entry, input_multiple_entry) + input_videos = self.__get_vid_paths(batch_mode, p['clip_single'], p['clip_list']) if len(input_videos) == 0: # exit if no paths found return with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: for video_path in input_videos: output_directory = self.__get_output_dir( - self.output_subdir_clip_entry.get(), batch_mode, - output_entry, video_path, input_multiple_entry) - time_start = "00:00:00" if batch_mode else str(self.clip_time_start_entry.get()) - time_end = "99:99:99" if batch_mode else str(self.clip_time_end_entry.get()) + p['output_subdir_clip'], batch_mode, + p['clip_output'], video_path, p['clip_list']) + time_start = "00:00:00" if batch_mode else p['clip_time_start'] + time_end = "99:99:99" if batch_mode else p['clip_time_end'] executor.submit(self.__extract_clips, str(video_path), time_start, time_end, max_length, - self.split_at_cuts.get(), bool(self.clip_bordercrop_entry.get()), + p['split_cuts'], p['clip_bordercrop'], crop_variation, target_fps, output_directory) if batch_mode: - self.__update_status(f'Clip extraction from all videos in "{input_multiple_entry}" complete') + self.__update_status(f'Clip extraction from all videos in "{p["clip_list"]}" complete') else: - self.__update_status(f'Clip extraction from "{input_single_entry}" complete') + self.__update_status(f'Clip extraction from "{p["clip_single"]}" complete') def __extract_clips(self, video_path: str, timestamp_min: str, timestamp_max: str, max_length: float, split_at_cuts: bool, remove_borders: bool, crop_variation: float, target_fps: float, output_dir: str): @@ -614,11 +326,9 @@ def __save_clip(self, scene: tuple[int, int], video_path: str, target_fps: float preview = Image.fromarray( cv2.cvtColor(frame[y1+y2:y1+y2+h2, x1+x2:x1+x2+w2], cv2.COLOR_BGR2RGB)) preview.thumbnail((150, 150)) - self.preview_image.configure(light_image=preview, size=preview.size) #truncate filename of long files so UI doesn't shift around filename_truncated = basename + ext if len(basename) < 20 else basename[:18] + ".." + ext - self.preview_image_label.configure( - text=f'{filename_truncated}\nFrames: {scene[0]}-{scene[1]}\nSize: {w2}x{h2}') + self.view.update_preview(preview, f'{filename_truncated}\nFrames: {scene[0]}-{scene[1]}\nSize: {w2}x{h2}') except Exception: pass video.release() @@ -701,22 +411,20 @@ def __write_clip_av(video_path: str, output_path: str, scene: tuple[int, int], for pkt in out_video.encode(): output_container.mux(pkt) - def __extract_images_button(self, batch_mode: bool): + def extract_images_button(self, batch_mode: bool): self.__run_in_thread(self.__extract_images_multi, batch_mode) - def __extract_images_multi(self, batch_mode : bool): - if not pathlib.Path(self.image_output_entry.get()).is_dir() or self.image_output_entry.get() == "": + def __extract_images_multi(self, batch_mode: bool): + p = self.args + if not pathlib.Path(p['image_output']).is_dir() or p['image_output'] == "": self.__update_status("Invalid output directory!") return # validate numeric inputs try: - capture_rate = float(self.capture_rate_entry.get()) - blur_threshold = float(self.blur_threshold_entry.get()) - crop_variation = float(self.image_crop_entry.get()) - input_single_entry = self.image_single_entry.get() - input_multiple_entry = self.image_list_entry.get() - output_entry = self.image_output_entry.get() + capture_rate = float(p['capture_rate']) + blur_threshold = float(p['blur_threshold']) + crop_variation = float(p['image_crop']) except ValueError: self.__update_status("Invalid numeric input for Images/sec, Blur Removal, or Crop Variation.") return @@ -730,25 +438,25 @@ def __extract_images_multi(self, batch_mode : bool): self.__update_status("Crop Variation must be between 0.0 and 1.0.") return - input_videos = self.__get_vid_paths(batch_mode, input_single_entry, input_multiple_entry) + input_videos = self.__get_vid_paths(batch_mode, p['image_single'], p['image_list']) if not input_videos: return with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: for video_path in input_videos: output_directory = self.__get_output_dir( - self.output_subdir_img_entry.get(), batch_mode, - output_entry, video_path, input_multiple_entry) - time_start = "00:00:00" if batch_mode else str(self.image_time_start_entry.get()) - time_end = "99:99:99" if batch_mode else str(self.image_time_end_entry.get()) + p['output_subdir_img'], batch_mode, + p['image_output'], video_path, p['image_list']) + time_start = "00:00:00" if batch_mode else p['image_time_start'] + time_end = "99:99:99" if batch_mode else p['image_time_end'] executor.submit(self.__save_frames, str(video_path), time_start, time_end, capture_rate, - blur_threshold, self.image_bordercrop.get(), + blur_threshold, p['image_bordercrop'], crop_variation, output_directory) if batch_mode: - self.__update_status(f'Image extraction from all videos in {input_multiple_entry} complete') + self.__update_status(f'Image extraction from all videos in {p["image_list"]} complete') else: - self.__update_status(f'Image extraction from "{input_single_entry}" complete') + self.__update_status(f'Image extraction from "{p["image_single"]}" complete') def __save_frames(self, video_path: str, timestamp_min: str, timestamp_max: str, capture_rate: float, blur_threshold: float, remove_borders: bool, crop_variation: float, output_dir: str): @@ -821,32 +529,26 @@ def __save_frames(self, video_path: str, timestamp_min: str, timestamp_max: str, cv2.cvtColor(frame_cropped[y2:y2+h2, x2:x2+w2], cv2.COLOR_BGR2RGB)) preview.thumbnail((150, 150)) filename_truncated = basename + ext if len(basename) < 20 else basename[:17] + "..." + ext - self.preview_image.configure(light_image=preview, size=preview.size) - self.preview_image_label.configure(text=f'{filename_truncated}\nFrame: {f[0]}\nSize: {w2}x{h2}') + self.view.update_preview(preview, f'{filename_truncated}\nFrame: {f[0]}\nSize: {w2}x{h2}') except Exception: pass # preview update is non-critical cv2.imwrite(filename, frame_cropped[y2:y2+h2, x2:x2+w2]) video.release() - def __download_button(self, batch_mode: bool): + def download_button(self, batch_mode: bool): self.__run_in_thread(self.__download_multi, batch_mode) - def __update_status(self, status_text: str): - print(status_text) - self.status_label.configure(state="normal") - self.status_label.insert(index="end", text=status_text + "\n") - self.status_label.configure(state="disabled") - def __download_multi(self, batch_mode: bool): - if not pathlib.Path(self.download_output_entry.get()).is_dir() or self.download_output_entry.get() == "": + p = self.args + if not pathlib.Path(p['download_output']).is_dir() or p['download_output'] == "": self.__update_status("Invalid output directory!") return if not batch_mode: - ydl_urls = [self.download_link_entry.get()] + ydl_urls = [p['download_link']] elif batch_mode: - ydl_path = pathlib.Path(self.download_list_entry.get()) + ydl_path = pathlib.Path(p['download_list']) if ydl_path.is_file() and ydl_path.suffix.lower() == ".txt": with open(ydl_path) as file: ydl_urls = file.readlines() @@ -857,8 +559,8 @@ def __download_multi(self, batch_mode: bool): with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: for url in ydl_urls: executor.submit(self.__download_video, - url.strip(), self.download_output_entry.get(), - self.download_args_entry.get("0.0", ctk.END)) + url.strip(), p['download_output'], + p['download_args']) self.__update_status(f'Completed {len(ydl_urls)} downloads.') diff --git a/modules/util/path_util.py b/modules/util/path_util.py index 94f802423..7d25c8e8c 100644 --- a/modules/util/path_util.py +++ b/modules/util/path_util.py @@ -1,5 +1,6 @@ import json import os.path +from pathlib import Path from typing import Any @@ -52,3 +53,8 @@ def supported_video_extensions() -> set[str]: def is_supported_video_extension(extension: str) -> bool: return extension.lower() in SUPPORTED_VIDEO_EXTENSIONS + + +def json_path_modifier(x: str | Path) -> Path: + x = Path(x).absolute() + return x.parent if x.suffix == ".json" else x diff --git a/modules/util/ui/CtkUIState.py b/modules/util/ui/CtkUIState.py new file mode 100644 index 000000000..0096ca3e6 --- /dev/null +++ b/modules/util/ui/CtkUIState.py @@ -0,0 +1,23 @@ +import tkinter as tk +from typing import Any + +from modules.util.ui.UIState import BaseUIState + + +class CtkUIState(BaseUIState): + def __init__(self, master, obj): + self.master = master + super().__init__(obj) + + def _make_str_var(self, initial_value: Any): + var = tk.StringVar(master=self.master) + var.set(initial_value) + return var + + def _make_bool_var(self, initial_value: Any): + var = tk.BooleanVar(master=self.master) + var.set(initial_value) + return var + + def _make_nested_state(self, obj: Any) -> "CtkUIState": + return CtkUIState(self.master, obj) diff --git a/modules/util/ui/UIState.py b/modules/util/ui/UIState.py index 8b13d23f7..73d73ed34 100644 --- a/modules/util/ui/UIState.py +++ b/modules/util/ui/UIState.py @@ -1,4 +1,4 @@ -import tkinter as tk +from abc import ABC, abstractmethod from collections.abc import Callable from dataclasses import dataclass from enum import Enum @@ -8,13 +8,12 @@ from modules.util.type_util import issubclass_safe -class UIState: +class BaseUIState(ABC): __vars: dict[str, Any] __var_traces: dict[str, dict[int, Callable[[], None]]] __latest_var_trace_id: int - def __init__(self, master, obj): - self.master = master + def __init__(self, obj): self.obj = obj self.__var_types: dict[str, type] = {} @@ -25,13 +24,24 @@ def __init__(self, master, obj): self.__var_traces = {name: {} for name in self.__vars} self.__latest_var_trace_id = 0 + @abstractmethod + def _make_str_var(self, initial_value: Any): + pass + + @abstractmethod + def _make_bool_var(self, initial_value: Any): + pass + + @abstractmethod + def _make_nested_state(self, obj: Any) -> "BaseUIState": + pass + def update(self, obj): self.obj = obj self.__set_vars(obj) def get_var(self, name): split_name = name.split('.') - if len(split_name) == 1: return self.__vars[split_name[0]] else: @@ -72,7 +82,6 @@ def update(_0, _1, _2): else: setattr(obj, name, string_var) self.__call_var_traces(name) - return update def __set_enum_var(self, obj, is_dict, name, var, var_type, nullable): @@ -92,7 +101,6 @@ def update(_0, _1, _2): else: setattr(obj, name, var_type[string_var]) self.__call_var_traces(name) - return update def __set_bool_var(self, obj, is_dict, name, var): @@ -104,7 +112,6 @@ def update(_0, _1, _2): def update(_0, _1, _2): setattr(obj, name, var.get()) self.__call_var_traces(name) - return update def __set_int_var(self, obj, is_dict, name, var, nullable): @@ -138,7 +145,6 @@ def update(_0, _1, _2): except ValueError: setattr(obj, name, None) self.__call_var_traces(name) - return update def __set_float_var(self, obj, is_dict, name, var, nullable): @@ -172,12 +178,10 @@ def update(_0, _1, _2): except ValueError: setattr(obj, name, None) self.__call_var_traces(name) - return update def __create_vars(self, obj): new_vars = {} - is_dict = isinstance(obj, dict) is_config = isinstance(obj, BaseConfig) @@ -190,61 +194,48 @@ def __create_vars(self, obj): obj_var = getattr(obj, name) if issubclass_safe(var_type, BaseConfig): - var = UIState(self.master, obj_var) - new_vars[name] = var + new_vars[name] = self._make_nested_state(obj_var) elif var_type is str: - var = tk.StringVar(master=self.master) - var.set("" if obj_var is None else obj_var) + var = self._make_str_var("" if obj_var is None else obj_var) var.trace_add("write", self.__set_str_var(obj, is_dict, name, var, obj.nullables[name])) new_vars[name] = var elif issubclass_safe(var_type, Enum): - var = tk.StringVar(master=self.master) - var.set("" if obj_var is None else str(obj_var)) + var = self._make_str_var("" if obj_var is None else str(obj_var)) var.trace_add("write", self.__set_enum_var(obj, is_dict, name, var, var_type, obj.nullables[name])) new_vars[name] = var elif var_type is bool: - var = tk.BooleanVar(master=self.master) - var.set(obj_var or False) + var = self._make_bool_var(obj_var or False) var.trace_add("write", self.__set_bool_var(obj, is_dict, name, var)) new_vars[name] = var elif var_type is int: - var = tk.StringVar(master=self.master) - var.set("" if obj_var is None else str(obj_var)) + var = self._make_str_var("" if obj_var is None else str(obj_var)) var.trace_add("write", self.__set_int_var(obj, is_dict, name, var, obj.nullables[name])) new_vars[name] = var elif var_type is float: - var = tk.StringVar(master=self.master) - var.set("" if obj_var is None else str(obj_var)) + var = self._make_str_var("" if obj_var is None else str(obj_var)) var.trace_add("write", self.__set_float_var(obj, is_dict, name, var, obj.nullables[name])) new_vars[name] = var else: iterable = obj.items() if is_dict else vars(obj).items() - for name, obj_var in iterable: - if isinstance(obj_var, str): - var = tk.StringVar(master=self.master) - var.set(obj_var) + var = self._make_str_var(obj_var) var.trace_add("write", self.__set_str_var(obj, is_dict, name, var, False)) new_vars[name] = var elif isinstance(obj_var, Enum): - var = tk.StringVar(master=self.master) - var.set(str(obj_var)) + var = self._make_str_var(str(obj_var)) var.trace_add("write", self.__set_enum_var(obj, is_dict, name, var, type(obj_var), False)) new_vars[name] = var elif isinstance(obj_var, bool): - var = tk.BooleanVar(master=self.master) - var.set(obj_var) + var = self._make_bool_var(obj_var) var.trace_add("write", self.__set_bool_var(obj, is_dict, name, var)) new_vars[name] = var elif isinstance(obj_var, int): - var = tk.StringVar(master=self.master) - var.set(str(obj_var)) + var = self._make_str_var(str(obj_var)) var.trace_add("write", self.__set_int_var(obj, is_dict, name, var, False)) new_vars[name] = var elif isinstance(obj_var, float): - var = tk.StringVar(master=self.master) - var.set(str(obj_var)) + var = self._make_str_var(str(obj_var)) var.trace_add("write", self.__set_float_var(obj, is_dict, name, var, False)) new_vars[name] = var @@ -253,7 +244,6 @@ def __create_vars(self, obj): def __set_vars(self, obj): is_dict = isinstance(obj, dict) is_config = isinstance(obj, BaseConfig) - iterable = obj.items() if is_dict else vars(obj).items() if is_config: for name, var_type in obj.types.items(): @@ -274,6 +264,7 @@ def __set_vars(self, obj): var = self.__vars[name] var.set("" if obj_var is None else str(obj_var)) else: + iterable = obj.items() if is_dict else vars(obj).items() for name, obj_var in iterable: if isinstance(obj_var, str): var = self.__vars[name] @@ -288,27 +279,26 @@ def __set_vars(self, obj): var = self.__vars[name] var.set(str(obj_var)) - # metadata api + @dataclass(frozen=True) + class VarMeta: + type: type | None + nullable: bool + default: Any + def _resolve_state_and_leaf(self, name: str): parts = name.split('.') - state: UIState = self + state: BaseUIState = self for part in parts[:-1]: state = state.get_var(part) - if not isinstance(state, UIState): + if not isinstance(state, BaseUIState): return None, None return state, parts[-1] - @dataclass(frozen=True) - class VarMeta: - type: type | None - nullable: bool - default: Any - - def get_field_metadata(self, name: str) -> "UIState.VarMeta": + def get_field_metadata(self, name: str) -> "BaseUIState.VarMeta": state, leaf = self._resolve_state_and_leaf(name) if state is None: - return UIState.VarMeta(None, False, None) - return UIState.VarMeta( + return BaseUIState.VarMeta(None, False, None) + return BaseUIState.VarMeta( state.__var_types.get(leaf), state.__var_nullables.get(leaf, False), state.__var_defaults.get(leaf, None), diff --git a/modules/util/ui/components.py b/modules/util/ui/ctk_components.py similarity index 88% rename from modules/util/ui/components.py rename to modules/util/ui/ctk_components.py index dd7b89719..e462f72a1 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/ctk_components.py @@ -8,9 +8,9 @@ from modules.util.enum.PathIOType import PathIOType from modules.util.enum.TimeUnit import TimeUnit from modules.util.path_util import supported_image_extensions +from modules.util.ui.ctk_validation import DEFAULT_MAX_UNDO, FieldValidator, PathValidator +from modules.util.ui.CtkUIState import CtkUIState from modules.util.ui.ToolTip import ToolTip -from modules.util.ui.UIState import UIState -from modules.util.ui.validation import DEFAULT_MAX_UNDO, FieldValidator, PathValidator import customtkinter as ctk from customtkinter.windows.widgets.scaling import CTkScalingBaseClass @@ -34,11 +34,13 @@ def app_title(master, row, column): 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): +def label(master, row, column, text, pad=PAD, tooltip=None, wide_tooltip=False, wraplength=0, underline=False): 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) + if underline: + component.configure(font=ctk.CTkFont(underline=True)) return component @@ -46,7 +48,7 @@ def entry( master, row, column, - ui_state: UIState, + ui_state: CtkUIState, var_name: str, command: Callable[[], None] | None = None, tooltip: str = "", @@ -108,13 +110,8 @@ def new_destroy(): return component -def json_path_modifier(x: str | Path) -> Path: - x = Path(x).absolute() - return x.parent if x.suffix == ".json" else x - - def path_entry( - master, row, column, ui_state: UIState, var_name: str, + master, row, column, ui_state: CtkUIState, var_name: str, *, mode: Literal["file", "dir"] = "file", io_type: PathIOType = PathIOType.INPUT, @@ -124,9 +121,10 @@ def path_entry( command: Callable[[str], None] | None = None, extra_validate: Callable[[str], str | None] | None = None, required: bool = False, + columnspan: int = 1, ): frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") + frame.grid(row=row, column=column, padx=0, pady=0, sticky="new", columnspan=columnspan) frame.grid_columnconfigure(0, weight=1) @@ -216,7 +214,7 @@ def _frame_destroy(): return frame -def time_entry(master, row, column, ui_state: UIState, var_name: str, unit_var_name, supports_time_units: bool = True): +def time_entry(master, row, column, ui_state: CtkUIState, 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") @@ -239,7 +237,7 @@ def time_entry(master, row, column, ui_state: UIState, var_name: str, unit_var_n return frame -def layer_filter_entry(master, row, column, ui_state: UIState, preset_var_name: str, preset_label: str, preset_tooltip: str, presets, entry_var_name, entry_tooltip: str, regex_var_name, regex_tooltip: str, frame_color=None): +def layer_filter_entry(master, row, column, ui_state: CtkUIState, preset_var_name: str, preset_label: str, preset_tooltip: str, presets, entry_var_name, entry_tooltip: str, regex_var_name, regex_tooltip: str, frame_color=None): frame = ctk.CTkFrame(master=master, corner_radius=5, fg_color=frame_color) frame.grid(row=row, column=column, padx=5, pady=5, sticky="nsew") frame.grid_columnconfigure(0, weight=1) @@ -353,6 +351,15 @@ def icon_button(master, row, column, text, command): return component +def colored_icon_button(master, row, column, text, fg_color, command, padx=0): + component = ctk.CTkButton( + master=master, width=20, height=20, text=text, + corner_radius=2, fg_color=fg_color, command=command, + ) + component.grid(row=row, column=column, padx=padx) + return component + + def button(master, row, column, text, command, tooltip=None, **kwargs): # Pop grid-specific parameters from kwargs, using PAD as the default if not provided. padx = kwargs.pop('padx', PAD) @@ -365,7 +372,7 @@ def button(master, row, column, text, command, tooltip=None, **kwargs): return component -def options(master, row, column, values, ui_state: UIState, var_name: str, command: Callable[[str], None] | None = None): +def options(master, row, column, values, ui_state: CtkUIState, var_name: str, command: Callable[[str], None] | None = None): 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") @@ -385,7 +392,7 @@ def destroy(self): return component -def options_adv(master, row, column, values, ui_state: UIState, var_name: str, +def options_adv(master, row, column, values, ui_state: CtkUIState, var_name: str, command: Callable[[str], None] | None = None, adv_command: Callable[[], None] | None = None): frame = ctk.CTkFrame(master, fg_color="transparent") frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") @@ -417,7 +424,7 @@ def destroy(self): return frame, {'component': component, 'button_component': button_component} -def options_kv(master, row, column, values: list[tuple[str, Any]], ui_state: UIState, var_name: str, +def options_kv(master, row, column, values: list[tuple[str, Any]], ui_state: CtkUIState, var_name: str, command: Callable[[Any], None] | None = None): var = ui_state.get_var(var_name) keys = [key for key, value in values] @@ -475,16 +482,19 @@ def switch( master, row, column, - ui_state: UIState, + ui_state: CtkUIState, var_name: str, command: Callable[[], None] | None = None, text: str = "", + width: int | None = None, ): var = ui_state.get_var(var_name) if command: trace_id = ui_state.add_var_trace(var_name, command) component = ctk.CTkSwitch(master, variable=var, text=text, command=command) + if width is not None: + component.configure(width=width) component.grid(row=row, column=column, padx=PAD, pady=(PAD, PAD), sticky="new") def create_destroy(component): @@ -545,3 +555,34 @@ def set_2(value, max_value): description_2_component.configure(text=f"{value}/{max_value}") return set_1, set_2 + + +def section_frame(master, row: int, col: int = 0): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=col, padx=PAD // 2, pady=PAD // 2, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + return frame + + +def inline_frame(master, row: int, col: int, columnspan: int = 1): + frame = ctk.CTkFrame(master, fg_color="transparent") + frame.grid(row=row, column=col, columnspan=columnspan, sticky="ew", padx=0, pady=0) + return frame + + +def set_widget_enabled(widget, enabled: bool) -> None: + state = "normal" if enabled else "disabled" + if isinstance(widget, ctk.CTkFrame): + for child in widget.children.values(): + with contextlib.suppress(Exception): + child.configure(state=state) + else: + widget.configure(state=state) + + +def set_label_text(label, text: str) -> None: + label.configure(text=str(text)) + + +def call_after(widget, delay_ms: int, func) -> None: + widget.after(delay_ms, func) diff --git a/modules/util/ui/ctk_validation.py b/modules/util/ui/ctk_validation.py index 9f44de8eb..5347ba6d4 100644 --- a/modules/util/ui/ctk_validation.py +++ b/modules/util/ui/ctk_validation.py @@ -1,18 +1,21 @@ from __future__ import annotations import contextlib -import os -import re -import sys import tkinter as tk -from collections import deque from collections.abc import Callable -from pathlib import PurePosixPath, PureWindowsPath from typing import TYPE_CHECKING, Any -from urllib.parse import urlparse -from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.PathIOType import PathIOType +from modules.util.ui.validation import ( + DEBOUNCE_TYPING_MS, + DEFAULT_MAX_UNDO, + ERROR_BORDER_COLOR, + UNDO_DEBOUNCE_MS, + BaseFieldValidator, + UndoHistory, + _active_validators, + _validate_path_field, +) if TYPE_CHECKING: from modules.util.ui.UIState import UIState @@ -20,159 +23,6 @@ import customtkinter as ctk -DEBOUNCE_TYPING_MS = 250 -UNDO_DEBOUNCE_MS = 500 -ERROR_BORDER_COLOR = "#dc3545" - -_active_validators: set[FieldValidator] = set() - -TRAILING_SLASH_RE = re.compile(r"[\\/]$") -ENDS_WITH_EXT = re.compile(r"\.[A-Za-z0-9]+$") -HUGGINGFACE_REPO_RE = re.compile(r"^[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+$") - -_INVALID_CHARS = {chr(c) for c in range(32)} -_IS_WINDOWS = sys.platform == "win32" -if _IS_WINDOWS: - _INVALID_CHARS |= set('<>"|?*') - - -def _is_huggingface_repo_or_file(value: str) -> bool: - trimmed = value.strip() - - if trimmed.startswith("https://"): - parsed = urlparse(trimmed) - if parsed.netloc not in {"huggingface.co", "huggingface.com"}: - return False - parts = parsed.path.strip("/").split("/") - if len(parts) >= 5 and parts[2] in {"resolve", "blob"}: - return bool(ENDS_WITH_EXT.search(parts[-1])) - return False - - if len(trimmed) > 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 len(trimmed) >= 2 and trimmed[1] == ":" and trimmed[0].isalpha(): - return False - if trimmed.count("/") != 1: - return False - - return bool(HUGGINGFACE_REPO_RE.match(trimmed)) - - -def _has_invalid_chars(value: str) -> bool: - return bool(_INVALID_CHARS.intersection(value)) - - -def _check_overwrite(path: str, *, is_dir: bool, prevent: bool) -> str | None: - if not prevent: - return None - abs_path = os.path.abspath(path) - if is_dir and os.path.isdir(abs_path): - return "Output folder already exists (overwrite prevented)" - if not is_dir and os.path.isfile(abs_path): - return "Output file already exists (overwrite prevented)" - return None - - -def validate_path( - value: str, - io_type: PathIOType = PathIOType.INPUT, - *, - prevent_overwrites: bool = False, - output_format: str | None = None, -) -> str | None: - """Return an error string if *value* is an invalid path, else ``None``.""" - trimmed = value.strip() - - if not trimmed: - return "Path is empty" - if TRAILING_SLASH_RE.search(trimmed): - return "Path must not end with a slash" - if _has_invalid_chars(trimmed): - return "Path contains invalid characters" - - if trimmed.startswith("cloud:"): - cloud_path = trimmed[6:] - if not cloud_path: - return "Cloud path is empty" - if cloud_path.startswith(("http://", "https://")): - return "Cloud path cannot be a URL" - if "\\" in cloud_path: - return "Cloud path must use forward slashes (/)" - return None - - if io_type == PathIOType.INPUT and _is_huggingface_repo_or_file(trimmed): - return None - - if io_type == PathIOType.INPUT: - if not os.path.exists(os.path.abspath(trimmed)): - return "Input path does not exist" - - if io_type in (PathIOType.OUTPUT, PathIOType.MODEL): - if not os.path.isdir(os.path.dirname(os.path.abspath(trimmed))): - return "Parent folder does not exist" - - if io_type == PathIOType.MODEL and output_format is not None: - if output_format == "DIFFUSERS": - if ENDS_WITH_EXT.search(trimmed): - return "Diffusers output must be a directory path, not a file" - return _check_overwrite(trimmed, is_dir=True, prevent=prevent_overwrites) - - try: - expected_ext = ModelFormat[output_format].file_extension() - except KeyError: - expected_ext = "" - - if expected_ext: - suffix = (PureWindowsPath(trimmed) if _IS_WINDOWS else PurePosixPath(trimmed)).suffix.lower() - if suffix != expected_ext: - return f"Extension must be '{expected_ext}' for {output_format} format" - return _check_overwrite(trimmed, is_dir=False, prevent=prevent_overwrites) - - if io_type == PathIOType.OUTPUT: - return _check_overwrite(trimmed, is_dir=False, prevent=prevent_overwrites) - - return None - -DEFAULT_MAX_UNDO = 20 - - -class UndoHistory: - def __init__(self, max_size: int = DEFAULT_MAX_UNDO): - self._stack: deque[str] = deque(maxlen=max_size) - self._redo_stack: list[str] = [] - - def push(self, value: str): - if self._stack and self._stack[-1] == value: - return - self._stack.append(value) - self._redo_stack.clear() - - def undo(self, current: str) -> str | None: - if not self._stack: - return None - top = self._stack[-1] - if top == current and len(self._stack) > 1: - self._redo_stack.append(self._stack.pop()) - return self._stack[-1] - elif top != current: - self._redo_stack.append(current) - return top - return None - - def redo(self) -> str | None: - if not self._redo_stack: - return None - value = self._redo_stack.pop() - self._stack.append(value) - return value - - class DebounceTimer: def __init__(self, widget, delay_ms: int, callback: Callable[..., Any]): self.widget = widget @@ -199,7 +49,7 @@ def cancel(self): self._after_id = None -class FieldValidator: +class FieldValidator(BaseFieldValidator): def __init__( self, component: ctk.CTkEntry, @@ -210,12 +60,9 @@ def __init__( extra_validate: Callable[[str], str | None] | None = None, required: bool = False, ): + super().__init__(ui_state, var_name, extra_validate, required) self.component = component self.var = var - self.ui_state = ui_state - self.var_name = var_name - self._extra_validate = extra_validate - self._required = required try: self._original_border_color = component.cget("border_color") @@ -227,7 +74,6 @@ def __init__( self._real_var_trace_name: str | None = None self._syncing = False self._touched = False - self._bound = False self._debounce: DebounceTimer | None = None self._undo_debounce: DebounceTimer | None = None @@ -309,59 +155,12 @@ def _commit(self) -> None: self.var.set(shadow_val) self._syncing = False - def validate(self, value: str) -> str | None: - """Return an error string if *value* is invalid, else None.""" - meta = self.ui_state.get_field_metadata(self.var_name) - declared_type = meta.type - nullable = meta.nullable - default_val = meta.default - - if value == "": - if self._required: - return "Value required" - if nullable: - return None - if declared_type is str: - if default_val == "": - return None - return "Value required" - return None - - try: - if declared_type is int: - v = int(value) - if v < 0: - return "Value must be non-negative" - elif declared_type is float: - v = float(value) - if v < 0: - return "Value must be non-negative" - elif declared_type is bool: - if value.lower() not in ("true", "false", "0", "1"): - return "Invalid bool" - except ValueError: - return "Invalid value" - - if self._extra_validate is not None: - return self._extra_validate(value) - - return None - def _apply_error(self) -> None: self.component.configure(border_color=ERROR_BORDER_COLOR) def _clear_error(self) -> None: self.component.configure(border_color=self._original_border_color) - def _validate_and_style(self, value: str) -> bool: - error = self.validate(value) - if error is None: - self._clear_error() - return True - else: - self._apply_error() - return False - def _on_shadow_write(self, *_args) -> None: if self._syncing: return @@ -436,6 +235,21 @@ def _on_redo(self, _e=None) -> str: self._set_value(next_val) return "break" + def flush(self) -> str | None: + if self._debounce: + self._debounce.cancel() + + value = self._shadow_var.get() + error = self.validate(value) + + if error is not None: + self._apply_error() + else: + self._clear_error() + self._commit() + + return error + class PathValidator(FieldValidator): """FieldValidator with additional path-specific checks.""" @@ -454,48 +268,14 @@ def __init__( super().__init__(component, var, ui_state, var_name, max_undo=max_undo, extra_validate=extra_validate, required=required) self.io_type = io_type - def _get_var_safe(self, name: str) -> tk.Variable | None: - try: - return self.ui_state.get_var(name) - except (KeyError, AttributeError): - return None - def validate(self, value: str) -> str | None: base_err = super().validate(value) if base_err is not None: return base_err if value == "": return None - - prevent_var = self._get_var_safe("prevent_overwrites") - format_var = self._get_var_safe("output_model_format") - return validate_path( - value, - io_type=self.io_type, - prevent_overwrites=prevent_var.get() if prevent_var is not None else False, - output_format=format_var.get() if format_var is not None else None, - ) + return _validate_path_field(self.ui_state, self.io_type, value) def revalidate(self) -> None: if self.component.winfo_exists(): self._validate_and_style(self._shadow_var.get()) - - -def flush_and_validate_all() -> list[str]: - invalid: list[str] = [] - - for v in list(_active_validators): - if v._debounce: - v._debounce.cancel() - - value = v._shadow_var.get() - error = v.validate(value) - - if error is not None: - v._apply_error() - invalid.append(f"{v.var_name}: {error}") - else: - v._clear_error() - v._commit() - - return invalid diff --git a/modules/util/ui/validation.py b/modules/util/ui/validation.py index 9f44de8eb..abb131be3 100644 --- a/modules/util/ui/validation.py +++ b/modules/util/ui/validation.py @@ -1,31 +1,26 @@ from __future__ import annotations -import contextlib import os import re import sys -import tkinter as tk +from abc import ABC, abstractmethod from collections import deque from collections.abc import Callable from pathlib import PurePosixPath, PureWindowsPath -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING from urllib.parse import urlparse from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.PathIOType import PathIOType if TYPE_CHECKING: - from modules.util.ui.UIState import UIState - - import customtkinter as ctk + from modules.util.ui.UIState import BaseUIState DEBOUNCE_TYPING_MS = 250 UNDO_DEBOUNCE_MS = 500 ERROR_BORDER_COLOR = "#dc3545" -_active_validators: set[FieldValidator] = set() - TRAILING_SLASH_RE = re.compile(r"[\\/]$") ENDS_WITH_EXT = re.compile(r"\.[A-Za-z0-9]+$") HUGGINGFACE_REPO_RE = re.compile(r"^[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+$") @@ -173,141 +168,37 @@ def redo(self) -> str | None: return value -class DebounceTimer: - def __init__(self, widget, delay_ms: int, callback: Callable[..., Any]): - self.widget = widget - self.delay_ms = delay_ms - self.callback = callback - self._after_id: str | None = None - - def call(self, *args, **kwargs): - if self._after_id: - with contextlib.suppress(tk.TclError): - self.widget.after_cancel(self._after_id) - - def fire(): - self._after_id = None - self.callback(*args, **kwargs) - - with contextlib.suppress(tk.TclError): - self._after_id = self.widget.after(self.delay_ms, fire) - - def cancel(self): - if self._after_id: - with contextlib.suppress(tk.TclError): - self.widget.after_cancel(self._after_id) - self._after_id = None - - -class FieldValidator: +class BaseFieldValidator(ABC): def __init__( self, - component: ctk.CTkEntry, - var: tk.Variable, - ui_state: UIState, + ui_state: BaseUIState, var_name: str, - max_undo: int = DEFAULT_MAX_UNDO, extra_validate: Callable[[str], str | None] | None = None, required: bool = False, ): - self.component = component - self.var = var self.ui_state = ui_state self.var_name = var_name self._extra_validate = extra_validate self._required = required - - try: - self._original_border_color = component.cget("border_color") - except Exception: - self._original_border_color = "gray50" - - self._shadow_var = tk.StringVar(master=component) - self._shadow_trace_name: str | None = None - self._real_var_trace_name: str | None = None - self._syncing = False - self._touched = False self._bound = False - self._debounce: DebounceTimer | None = None - self._undo_debounce: DebounceTimer | None = None - self._undo = UndoHistory(max_undo) - - def attach(self) -> None: - self._shadow_var.set(self.var.get()) - self._swap_textvariable(self._shadow_var) - - self._debounce = DebounceTimer( - self.component, DEBOUNCE_TYPING_MS, self._on_debounce_fire - ) - self._undo_debounce = DebounceTimer( - self.component, UNDO_DEBOUNCE_MS, self._push_undo_snapshot - ) - - self._shadow_trace_name = self._shadow_var.trace_add("write", self._on_shadow_write) - self._real_var_trace_name = self.var.trace_add("write", self._on_real_var_write) - - self.component.bind("", self._on_focus_in) - self.component.bind("", self._on_user_input) - self.component.bind("<>", self._on_user_input) - self.component.bind("<>", self._on_user_input) - self.component.bind("", self._on_focus_out) - self.component.bind("", self._on_undo) - self.component.bind("", self._on_undo) - self.component.bind("", self._on_redo) - self.component.bind("", self._on_redo) - self.component.bind("", self._on_redo) - self.component.bind("", self._on_redo) - self.component.bind("", self._on_enter) - - self._bound = True - _active_validators.add(self) - - def detach(self) -> None: - if not self._bound: - return - self._bound = False - _active_validators.discard(self) - - self._commit() - - if self._debounce: - self._debounce.cancel() - if self._undo_debounce: - self._undo_debounce.cancel() - - if self._shadow_trace_name: - with contextlib.suppress(Exception): - self._shadow_var.trace_remove("write", self._shadow_trace_name) - self._shadow_trace_name = None - - if self._real_var_trace_name: - with contextlib.suppress(Exception): - self.var.trace_remove("write", self._real_var_trace_name) - self._real_var_trace_name = None - - self._swap_textvariable(self.var) - - def _swap_textvariable(self, new_var: tk.Variable) -> None: - comp = self.component - if comp._textvariable_callback_name: - with contextlib.suppress(Exception): - comp._textvariable.trace_remove("write", comp._textvariable_callback_name) # type: ignore[union-attr] - comp._textvariable_callback_name = "" + @abstractmethod + def _apply_error(self) -> None: + pass - comp.configure(textvariable=new_var) + @abstractmethod + def _clear_error(self) -> None: + pass - if new_var is not None: - comp._textvariable_callback_name = new_var.trace_add( - "write", comp._textvariable_callback - ) + @abstractmethod + def flush(self) -> str | None: + pass - def _commit(self) -> None: - shadow_val = self._shadow_var.get() - if shadow_val != self.var.get(): - self._syncing = True - self.var.set(shadow_val) - self._syncing = False + def _get_var_safe(self, name: str): + try: + return self.ui_state.get_var(name) + except (KeyError, AttributeError): + return None def validate(self, value: str) -> str | None: """Return an error string if *value* is invalid, else None.""" @@ -347,12 +238,6 @@ def validate(self, value: str) -> str | None: return None - def _apply_error(self) -> None: - self.component.configure(border_color=ERROR_BORDER_COLOR) - - def _clear_error(self) -> None: - self.component.configure(border_color=self._original_border_color) - def _validate_and_style(self, value: str) -> bool: error = self.validate(value) if error is None: @@ -362,140 +247,31 @@ def _validate_and_style(self, value: str) -> bool: self._apply_error() return False - def _on_shadow_write(self, *_args) -> None: - if self._syncing: - return - if not self._touched: - # external sync or initial set — commit immediately - self._commit() - if self._debounce: - self._debounce.cancel() - return - if self._debounce: - self._debounce.call() - if self._undo_debounce: - self._undo_debounce.call() - - def _on_real_var_write(self, *_args) -> None: - if self._syncing: - return - # external change (preset load, file dialog, etc) — sync to shadow var - self._syncing = True - self._shadow_var.set(self.var.get()) - self._syncing = False - self._validate_and_style(self._shadow_var.get()) - - def _push_undo_snapshot(self) -> None: - self._undo.push(self._shadow_var.get()) - - def _on_debounce_fire(self) -> None: - val = self._shadow_var.get() - if self._validate_and_style(val): - self._commit() - - def _on_focus_in(self, _e=None) -> None: - self._touched = False - self._undo.push(self._shadow_var.get()) - - def _on_user_input(self, _e=None) -> None: - self._touched = True - - def _on_focus_out(self, _e=None) -> None: - if self._debounce: - self._debounce.cancel() - if self._undo_debounce: - self._undo_debounce.cancel() - if self._touched: - if self._validate_and_style(self._shadow_var.get()): - self._commit() - self._undo.push(self._shadow_var.get()) - - def _on_enter(self, _e=None) -> None: - if self._debounce: - self._debounce.cancel() - if self._touched: - if self._validate_and_style(self._shadow_var.get()): - self._commit() - - def _set_value(self, value: str) -> None: - self._syncing = True - self._shadow_var.set(value) - self._syncing = False - if self._validate_and_style(value): - self._commit() - - def _on_undo(self, _e=None) -> str: - previous = self._undo.undo(self._shadow_var.get()) - if previous is not None: - self._set_value(previous) - return "break" - - def _on_redo(self, _e=None) -> str: - next_val = self._undo.redo() - if next_val is not None: - self._set_value(next_val) - return "break" - - -class PathValidator(FieldValidator): - """FieldValidator with additional path-specific checks.""" - - def __init__( - self, - component: ctk.CTkEntry, - var: tk.Variable, - ui_state: UIState, - var_name: str, - io_type: PathIOType = PathIOType.INPUT, - max_undo: int = DEFAULT_MAX_UNDO, - extra_validate: Callable[[str], str | None] | None = None, - required: bool = False, - ): - super().__init__(component, var, ui_state, var_name, max_undo=max_undo, extra_validate=extra_validate, required=required) - self.io_type = io_type - - def _get_var_safe(self, name: str) -> tk.Variable | None: - try: - return self.ui_state.get_var(name) - except (KeyError, AttributeError): - return None - - def validate(self, value: str) -> str | None: - base_err = super().validate(value) - if base_err is not None: - return base_err - if value == "": - return None - - prevent_var = self._get_var_safe("prevent_overwrites") - format_var = self._get_var_safe("output_model_format") - return validate_path( - value, - io_type=self.io_type, - prevent_overwrites=prevent_var.get() if prevent_var is not None else False, - output_format=format_var.get() if format_var is not None else None, - ) - def revalidate(self) -> None: - if self.component.winfo_exists(): - self._validate_and_style(self._shadow_var.get()) +_active_validators: set[BaseFieldValidator] = set() def flush_and_validate_all() -> list[str]: invalid: list[str] = [] - for v in list(_active_validators): - if v._debounce: - v._debounce.cancel() - - value = v._shadow_var.get() - error = v.validate(value) - + error = v.flush() if error is not None: - v._apply_error() invalid.append(f"{v.var_name}: {error}") - else: - v._clear_error() - v._commit() - return invalid + + +def _validate_path_field(ui_state: BaseUIState, io_type: PathIOType, value: str) -> str | None: + try: + prevent_var = ui_state.get_var("prevent_overwrites") + except (KeyError, AttributeError): + prevent_var = None + try: + format_var = ui_state.get_var("output_model_format") + except (KeyError, AttributeError): + format_var = None + return validate_path( + value, + io_type=io_type, + prevent_overwrites=prevent_var.get() if prevent_var is not None else False, + output_format=format_var.get() if format_var is not None else None, + ) diff --git a/scripts/train_ui.py b/scripts/train_ui.py index 46ee8f1e6..562c73feb 100644 --- a/scripts/train_ui.py +++ b/scripts/train_ui.py @@ -2,11 +2,11 @@ script_imports() -from modules.ui.TrainUI import TrainUI +from modules.ui.CtkTrainUIView import CtkTrainUIView def main(): - ui = TrainUI() + ui = CtkTrainUIView() ui.mainloop() From ec2ee0a9547a24e5bb8d54e4884d2c48681e602c Mon Sep 17 00:00:00 2001 From: dxqb Date: Mon, 11 May 2026 20:09:48 +0200 Subject: [PATCH 03/21] refactor: replace manual entry+browse-button pairs with path_entry in VideoToolUI Removes abstract _create_browse_dir_button/_create_browse_file_button from BaseVideoToolUIView and uses the combined path_entry component instead, fixing widget alignment issues. Adds allow_video_files flag to path_entry. Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/BaseVideoToolUIView.py | 39 +++++++++---------------------- modules/ui/CtkVideoToolUIView.py | 24 ------------------- modules/util/ui/ctk_components.py | 7 +++++- 3 files changed, 17 insertions(+), 53 deletions(-) diff --git a/modules/ui/BaseVideoToolUIView.py b/modules/ui/BaseVideoToolUIView.py index b82ba5701..e60247585 100644 --- a/modules/ui/BaseVideoToolUIView.py +++ b/modules/ui/BaseVideoToolUIView.py @@ -1,8 +1,6 @@ import webbrowser from abc import ABC, abstractmethod -from modules.util.path_util import SUPPORTED_VIDEO_EXTENSIONS - class BaseVideoToolUIView(ABC): def __init__(self, components): @@ -12,9 +10,8 @@ def build_clip_extract_tab(self, frame, controller, ui_state): # single video self.components.label(frame, 0, 0, "Single Video", tooltip="Link to single video file to process.") - self.components.entry(frame, 0, 1, ui_state, "clip_single", width=190) - self._create_browse_file_button(frame, 0, ui_state, "clip_single", - [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))]) + self.components.path_entry(frame, 0, 1, ui_state, "clip_single", + mode="file", allow_model_files=False, allow_video_files=True) self.components.button(frame, 0, 2, "Extract Single", command=lambda: self._extract_clips(False, controller)) @@ -28,16 +25,14 @@ def build_clip_extract_tab(self, frame, controller, ui_state): # directory of videos self.components.label(frame, 2, 0, "Directory", tooltip="Path to directory with multiple videos to process, including in subdirectories.") - self.components.entry(frame, 2, 1, ui_state, "clip_list", width=190) - self._create_browse_dir_button(frame, 2, ui_state, "clip_list") + self.components.path_entry(frame, 2, 1, ui_state, "clip_list", mode="dir") self.components.button(frame, 2, 2, "Extract Directory", command=lambda: self._extract_clips(True, controller)) # output directory self.components.label(frame, 3, 0, "Output", tooltip="Path to folder where extracted clips will be saved.") - self.components.entry(frame, 3, 1, ui_state, "clip_output", width=190) - self._create_browse_dir_button(frame, 3, ui_state, "clip_output") + self.components.path_entry(frame, 3, 1, ui_state, "clip_output", mode="dir") # output to subdirectories self.components.label(frame, 4, 0, "Output to\nSubdirectories", @@ -76,9 +71,8 @@ def build_image_extract_tab(self, frame, controller, ui_state): # single video self.components.label(frame, 0, 0, "Single Video", tooltip="Link to single video file to process.") - self.components.entry(frame, 0, 1, ui_state, "image_single", width=190) - self._create_browse_file_button(frame, 0, ui_state, "image_single", - [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))]) + self.components.path_entry(frame, 0, 1, ui_state, "image_single", + mode="file", allow_model_files=False, allow_video_files=True) self.components.button(frame, 0, 2, "Extract Single", command=lambda: self._extract_images(False, controller)) @@ -92,16 +86,14 @@ def build_image_extract_tab(self, frame, controller, ui_state): # directory of videos self.components.label(frame, 2, 0, "Directory", tooltip="Path to directory with multiple videos to process, including in subdirectories.") - self.components.entry(frame, 2, 1, ui_state, "image_list", width=190) - self._create_browse_dir_button(frame, 2, ui_state, "image_list") + self.components.path_entry(frame, 2, 1, ui_state, "image_list", mode="dir") self.components.button(frame, 2, 2, "Extract Directory", command=lambda: self._extract_images(True, controller)) # output directory self.components.label(frame, 3, 0, "Output", tooltip="Path to folder where extracted images will be saved.") - self.components.entry(frame, 3, 1, ui_state, "image_output", width=190) - self._create_browse_dir_button(frame, 3, ui_state, "image_output") + self.components.path_entry(frame, 3, 1, ui_state, "image_output", mode="dir") # output to subdirectories self.components.label(frame, 4, 0, "Output to\nSubdirectories", @@ -143,16 +135,15 @@ def build_video_download_tab(self, frame, controller, ui_state): # link list self.components.label(frame, 1, 0, "Link List", tooltip="Path to txt file with list of links separated by newlines.") - self.components.entry(frame, 1, 1, ui_state, "download_list", width=190) - self._create_browse_file_button(frame, 1, ui_state, "download_list", [("Text file", ".txt")]) + self.components.path_entry(frame, 1, 1, ui_state, "download_list", + mode="file", allow_model_files=False) self.components.button(frame, 1, 2, "Download List", command=lambda: self._download(True, controller)) # output directory self.components.label(frame, 2, 0, "Output", tooltip="Path to folder where downloaded videos will be saved.") - self.components.entry(frame, 2, 1, ui_state, "download_output", width=190) - self._create_browse_dir_button(frame, 2, ui_state, "download_output") + self.components.path_entry(frame, 2, 1, ui_state, "download_output", mode="dir") # additional args self.components.label(frame, 3, 0, "Additional Args", @@ -166,14 +157,6 @@ def build_video_download_tab(self, frame, controller, ui_state): def _create_textbox(self, master, row, col, width, height, ui_state, var_name): pass - @abstractmethod - def _create_browse_dir_button(self, master, row, ui_state, var_name): - pass - - @abstractmethod - def _create_browse_file_button(self, master, row, ui_state, var_name, filetypes): - pass - @abstractmethod def update_status(self, status_text: str): pass diff --git a/modules/ui/CtkVideoToolUIView.py b/modules/ui/CtkVideoToolUIView.py index c272c891c..d6190772f 100644 --- a/modules/ui/CtkVideoToolUIView.py +++ b/modules/ui/CtkVideoToolUIView.py @@ -1,5 +1,3 @@ -from tkinter import filedialog - from modules.ui.BaseVideoToolUIView import BaseVideoToolUIView from modules.ui.VideoToolUIController import VideoToolUIController from modules.util.image_util import load_image @@ -91,28 +89,6 @@ def on_text_change(event=None): textbox.bind("", on_text_change) return textbox - def _create_browse_dir_button(self, master, row, ui_state, var_name): - def browse(): - path = filedialog.askdirectory() - if path: - ui_state.get_var(var_name).set(path) - self.focus_set() - - button = ctk.CTkButton(master, width=30, text="...", command=browse) - button.grid(row=row, column=1, sticky="e", padx=PAD, pady=PAD) - return button - - def _create_browse_file_button(self, master, row, ui_state, var_name, filetypes): - def browse(): - path = filedialog.askopenfilename(filetypes=filetypes) - if path: - ui_state.get_var(var_name).set(path) - self.focus_set() - - button = ctk.CTkButton(master, width=30, text="...", command=browse) - button.grid(row=row, column=1, sticky="e", padx=PAD, pady=PAD) - return button - def update_status(self, status_text: str): self.status_label.configure(state="normal") self.status_label.insert(index="end", text=status_text + "\n") diff --git a/modules/util/ui/ctk_components.py b/modules/util/ui/ctk_components.py index e462f72a1..ec5e0ecf9 100644 --- a/modules/util/ui/ctk_components.py +++ b/modules/util/ui/ctk_components.py @@ -7,7 +7,7 @@ from modules.util.enum.PathIOType import PathIOType from modules.util.enum.TimeUnit import TimeUnit -from modules.util.path_util import supported_image_extensions +from modules.util.path_util import supported_image_extensions, supported_video_extensions from modules.util.ui.ctk_validation import DEFAULT_MAX_UNDO, FieldValidator, PathValidator from modules.util.ui.CtkUIState import CtkUIState from modules.util.ui.ToolTip import ToolTip @@ -118,6 +118,7 @@ def path_entry( path_modifier: Callable[[str], str | Path] | None = None, allow_model_files: bool = True, allow_image_files: bool = False, + allow_video_files: bool = False, command: Callable[[str], None] | None = None, extra_validate: Callable[[str], str | None] | None = None, required: bool = False, @@ -181,6 +182,10 @@ def __open_dialog(): filetypes.extend([ ("Image", ' '.join([f"*.{x}" for x in supported_image_extensions()])), ]) + if allow_video_files: + filetypes.extend([ + ("Video", ' '.join(f"*{e}" for e in supported_video_extensions())), + ]) if use_save_dialog: chosen = filedialog.asksaveasfilename(filetypes=filetypes, initialdir=current_dir, From c3961ba48afeacbb86c059fa1aa97ffe16d0d7f8 Mon Sep 17 00:00:00 2001 From: dxqb Date: Fri, 15 May 2026 11:40:15 +0200 Subject: [PATCH 04/21] fix: declare undeclared self attrs in Base*View classes Base view concrete methods were accessing self.controller, self.ui_state, and toolkit-specific action methods that were only set by CTK subclass __init__ after the base __init__ call, with no enforcement in the base. - BaseTrainUIView: add controller/ui_state as constructor params; fix sync_cloud_secrets to use controller.train_config; add @abstractmethod for export_training, generate_debug_package, open_profiling_tool - CtkTrainUIView: reorder __init__ to create deps before base init call; replace self.train_config with self.controller.train_config - BaseCloudTabView: add controller as constructor param - CtkCloudTabView: pass controller to base __init__, drop redundant assignment - BaseCaptionUIView: add ABC + @abstractmethod for 6 action callbacks - BaseConceptTabView: remove concrete _update_filters() (accessed CTK vars); add ConceptConfig import; add concept: ConceptConfig param to BaseConceptWidgetView.__init__ - CtkConceptTabView: implement _update_filters(); pass concept to base init - BaseConceptWindowView: initialize bucket_ax/text_color/canvas to None - BaseTrainingTabView: replace callbacks dict with 6 @abstractmethod declarations; restore_optimizer_config(variable: str) matches controller - CtkTrainingTabView: implement all 6 abstract methods directly Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/BaseCaptionUIView.py | 21 +++- modules/ui/BaseCloudTabView.py | 3 +- modules/ui/BaseConceptTabView.py | 14 ++- modules/ui/BaseConceptWindowView.py | 3 + modules/ui/BaseTrainUIView.py | 15 ++- modules/ui/BaseTrainingTabView.py | 184 +++++++++++++++------------- modules/ui/CtkCloudTabView.py | 3 +- modules/ui/CtkConceptTabView.py | 10 +- modules/ui/CtkTrainUIView.py | 32 ++--- modules/ui/CtkTrainingTabView.py | 34 +++-- 10 files changed, 184 insertions(+), 135 deletions(-) diff --git a/modules/ui/BaseCaptionUIView.py b/modules/ui/BaseCaptionUIView.py index a6561da69..c53ca84d8 100644 --- a/modules/ui/BaseCaptionUIView.py +++ b/modules/ui/BaseCaptionUIView.py @@ -1,10 +1,29 @@ import platform +from abc import ABC, abstractmethod -class BaseCaptionUIView: +class BaseCaptionUIView(ABC): def __init__(self, components): self.components = components + @abstractmethod + def open_directory(self): pass + + @abstractmethod + def open_mask_window(self): pass + + @abstractmethod + def open_caption_window(self): pass + + @abstractmethod + def open_in_explorer(self): pass + + @abstractmethod + def draw_mask_editing_mode(self, *args): pass + + @abstractmethod + def fill_mask_editing_mode(self, *args): pass + def build_top_bar(self, frame, controller, ui_state): self.components.button(frame, 0, 0, "Open", self.open_directory, tooltip="open a new directory") diff --git a/modules/ui/BaseCloudTabView.py b/modules/ui/BaseCloudTabView.py index 80be17360..3f20b5674 100644 --- a/modules/ui/BaseCloudTabView.py +++ b/modules/ui/BaseCloudTabView.py @@ -7,8 +7,9 @@ class BaseCloudTabView(ABC): - def __init__(self, components): + def __init__(self, components, controller): self.components = components + self.controller = controller @property def reattach(self): diff --git a/modules/ui/BaseConceptTabView.py b/modules/ui/BaseConceptTabView.py index 077a10c24..4d2fca570 100644 --- a/modules/ui/BaseConceptTabView.py +++ b/modules/ui/BaseConceptTabView.py @@ -4,6 +4,7 @@ from modules.ui.BaseConfigListView import BaseConfigListView from modules.ui.ConceptWindowController import ConceptWindowController from modules.util import path_util +from modules.util.config.ConceptConfig import ConceptConfig from modules.util.enum.ConceptType import ConceptType from modules.util.image_util import load_image @@ -14,6 +15,11 @@ class BaseConceptTabView(BaseConfigListView): _FILTER_TYPES = ["ALL", "STANDARD", "VALIDATION", "PRIOR_PREDICTION"] + def __init__(self, search_var, filter_var, show_disabled_var): + self.search_var = search_var + self.filter_var = filter_var + self.show_disabled_var = show_disabled_var + def _element_matches_filters(self, element): if not self.filters.get("show_disabled", True): if hasattr(element, 'enabled') and not element.enabled: @@ -50,17 +56,13 @@ def _element_matches_filters(self, element): return True - def _update_filters(self): - self._create_element_list(search=self.search_var.get(), - type=self.filter_var.get(), - show_disabled=self.show_disabled_var.get()) - self._refresh_show_disabled_text() class BaseConceptWidgetView: - def __init__(self, components): + def __init__(self, components, concept: ConceptConfig): self.components = components + self.concept = concept def _get_display_name(self): if self.concept.name: diff --git a/modules/ui/BaseConceptWindowView.py b/modules/ui/BaseConceptWindowView.py index 0f851f5fb..0b94e50d1 100644 --- a/modules/ui/BaseConceptWindowView.py +++ b/modules/ui/BaseConceptWindowView.py @@ -9,6 +9,9 @@ class BaseConceptWindowView: def __init__(self, components): self.components = components + self.bucket_ax = None + self.text_color = None + self.canvas = None def build_general_tab(self, frame, controller, ui_state, text_ui_state): # name diff --git a/modules/ui/BaseTrainUIView.py b/modules/ui/BaseTrainUIView.py index 7acbf0005..9a443a342 100644 --- a/modules/ui/BaseTrainUIView.py +++ b/modules/ui/BaseTrainUIView.py @@ -9,8 +9,10 @@ class BaseTrainUIView(ABC): - def __init__(self, components): + def __init__(self, components, controller, ui_state): self.components = components + self.controller = controller + self.ui_state = ui_state # --- Abstract callbacks (controller calls into view) --- @@ -51,7 +53,7 @@ def show_window(self, window): pass def connect_window_closed(self, window, callback): pass def sync_cloud_secrets(self): - self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) + self.ui_state.get_var("secrets.cloud").update(self.controller.train_config.secrets.cloud) def start_training(self): self.controller.start_training() @@ -83,6 +85,15 @@ def open_sampling_tool(self): pass @abstractmethod def open_manual_sample_window(self): pass + @abstractmethod + def open_profiling_tool(self): pass + + @abstractmethod + def export_training(self): pass + + @abstractmethod + def generate_debug_package(self): pass + # --- Content builders (components calls; called by CTK view after frame creation) --- def build_bottom_bar_content(self, frame, status_frame, controller, ui_state): diff --git a/modules/ui/BaseTrainingTabView.py b/modules/ui/BaseTrainingTabView.py index 22a6af34e..947ca4f56 100644 --- a/modules/ui/BaseTrainingTabView.py +++ b/modules/ui/BaseTrainingTabView.py @@ -1,4 +1,4 @@ -from abc import ABC +from abc import ABC, abstractmethod from modules.util.enum.DataType import DataType from modules.util.enum.EMAMode import EMAMode @@ -16,232 +16,250 @@ class BaseTrainingTabView(ABC): def __init__(self, components): self.components = components - def build(self, column_0, column_1, column_2, controller, ui_state, callbacks: dict): + @abstractmethod + def restore_optimizer_config(self, variable: str): pass + + @abstractmethod + def open_optimizer_params(self): pass + + @abstractmethod + def restore_scheduler(self, variable): pass + + @abstractmethod + def open_scheduler_params(self): pass + + @abstractmethod + def open_offloading(self): pass + + @abstractmethod + def open_timestep_distribution(self): pass + + def build(self, column_0, column_1, column_2, controller, ui_state): model_type = controller.config.model_type if model_type.is_stable_diffusion(): - self.__setup_stable_diffusion_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_stable_diffusion_ui(column_0, column_1, column_2, controller, ui_state) if model_type.is_stable_diffusion_3(): - self.__setup_stable_diffusion_3_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_stable_diffusion_3_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_stable_diffusion_xl(): - self.__setup_stable_diffusion_xl_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_stable_diffusion_xl_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_wuerstchen(): - self.__setup_wuerstchen_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_wuerstchen_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_pixart(): - self.__setup_pixart_alpha_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_pixart_alpha_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_flux_1(): - self.__setup_flux_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_flux_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_flux_2(): - self.__setup_flux_2_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_flux_2_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_chroma(): - self.__setup_chroma_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_chroma_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_qwen(): - self.__setup_qwen_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_qwen_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_sana(): - self.__setup_sana_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_sana_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_hunyuan_video(): - self.__setup_hunyuan_video_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_hunyuan_video_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_hi_dream(): - self.__setup_hi_dream_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_hi_dream_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_z_image(): - self.__setup_z_image_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_z_image_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_ernie(): - self.__setup_ernie_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_ernie_ui(column_0, column_1, column_2, controller, ui_state) - def __setup_stable_diffusion_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_stable_diffusion_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_frame(column_0, 1, ui_state) self.__create_embedding_frame(column_0, 2, ui_state) - self.__create_base2_frame(column_1, 0, ui_state, callbacks, supports_circular_padding=True) + self.__create_base2_frame(column_1, 0, ui_state, supports_circular_padding=True) self.__create_unet_frame(column_1, 1, ui_state) - self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_generalized_offset_noise=True) + self.__create_noise_frame(column_1, 2, ui_state, supports_generalized_offset_noise=True) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_stable_diffusion_3_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_stable_diffusion_3_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_n_frame(column_0, 1, ui_state, i=1, supports_include=True) self.__create_text_encoder_n_frame(column_0, 2, ui_state, i=2, supports_include=True) self.__create_text_encoder_n_frame(column_0, 3, ui_state, i=3, supports_include=True) self.__create_embedding_frame(column_0, 4, ui_state) - self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_base2_frame(column_1, 0, ui_state) self.__create_transformer_frame(column_1, 1, ui_state) - self.__create_noise_frame(column_1, 2, ui_state, callbacks) + self.__create_noise_frame(column_1, 2, ui_state) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_stable_diffusion_xl_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_stable_diffusion_xl_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_n_frame(column_0, 1, ui_state, i=1) self.__create_text_encoder_n_frame(column_0, 2, ui_state, i=2) self.__create_embedding_frame(column_0, 3, ui_state) - self.__create_base2_frame(column_1, 0, ui_state, callbacks, supports_circular_padding=True) + self.__create_base2_frame(column_1, 0, ui_state, supports_circular_padding=True) self.__create_unet_frame(column_1, 1, ui_state) - self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_generalized_offset_noise=True) + self.__create_noise_frame(column_1, 2, ui_state, supports_generalized_offset_noise=True) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_wuerstchen_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_wuerstchen_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_frame(column_0, 1, ui_state) self.__create_embedding_frame(column_0, 2, ui_state) - self.__create_base2_frame(column_1, 0, ui_state, callbacks, supports_circular_padding=True) + self.__create_base2_frame(column_1, 0, ui_state, supports_circular_padding=True) self.__create_prior_frame(column_1, 1, ui_state) - self.__create_noise_frame(column_1, 2, ui_state, callbacks) + self.__create_noise_frame(column_1, 2, ui_state) self.__create_masked_frame(column_2, 0, ui_state) self.__create_loss_frame(column_2, 1, controller, ui_state) self.__create_layer_frame(column_2, 2, controller, ui_state) - def __setup_pixart_alpha_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_pixart_alpha_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_frame(column_0, 1, ui_state) self.__create_embedding_frame(column_0, 2, ui_state) - self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_base2_frame(column_1, 0, ui_state) self.__create_transformer_frame(column_1, 1, ui_state) - self.__create_noise_frame(column_1, 2, ui_state, callbacks) + self.__create_noise_frame(column_1, 2, ui_state) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state, supports_vb_loss=True) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_flux_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_flux_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_n_frame(column_0, 1, ui_state, i=1, supports_include=True) self.__create_text_encoder_n_frame(column_0, 2, ui_state, i=2, supports_include=True, supports_sequence_length=True) self.__create_embedding_frame(column_0, 4, ui_state) - self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_base2_frame(column_1, 0, ui_state) self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=True) - self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_dynamic_timestep_shifting=True) + self.__create_noise_frame(column_1, 2, ui_state, supports_dynamic_timestep_shifting=True) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_flux_2_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_flux_2_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_frame(column_0, 1, ui_state, supports_clip_skip=False, supports_training=False, supports_sequence_length=True) - self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_base2_frame(column_1, 0, ui_state) self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=True, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_dynamic_timestep_shifting=True) + self.__create_noise_frame(column_1, 2, ui_state, supports_dynamic_timestep_shifting=True) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_chroma_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_chroma_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_frame(column_0, 1, ui_state) self.__create_embedding_frame(column_0, 4, ui_state) - self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_base2_frame(column_1, 0, ui_state) self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, ui_state, callbacks) + self.__create_noise_frame(column_1, 2, ui_state) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_qwen_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_qwen_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_frame(column_0, 1, ui_state, supports_clip_skip=False) - self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_base2_frame(column_1, 0, ui_state) self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_dynamic_timestep_shifting=True) + self.__create_noise_frame(column_1, 2, ui_state, supports_dynamic_timestep_shifting=True) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_z_image_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_z_image_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_frame(column_0, 1, ui_state, supports_clip_skip=False, supports_training=False) - self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_base2_frame(column_1, 0, ui_state) self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_dynamic_timestep_shifting=True) + self.__create_noise_frame(column_1, 2, ui_state, supports_dynamic_timestep_shifting=True) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_ernie_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_ernie_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_frame(column_0, 1, ui_state, supports_clip_skip=False, supports_training=False) - self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_base2_frame(column_1, 0, ui_state) self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_dynamic_timestep_shifting=True) + self.__create_noise_frame(column_1, 2, ui_state, supports_dynamic_timestep_shifting=True) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_sana_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_sana_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_frame(column_0, 1, ui_state) self.__create_embedding_frame(column_0, 2, ui_state) - self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_base2_frame(column_1, 0, ui_state) self.__create_transformer_frame(column_1, 1, ui_state) - self.__create_noise_frame(column_1, 2, ui_state, callbacks) + self.__create_noise_frame(column_1, 2, ui_state) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_hunyuan_video_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_hunyuan_video_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_n_frame(column_0, 1, ui_state, i=1, supports_include=True) self.__create_text_encoder_n_frame(column_0, 2, ui_state, i=2, supports_include=True) self.__create_embedding_frame(column_0, 4, ui_state) - self.__create_base2_frame(column_1, 0, ui_state, callbacks, video_training_enabled=True) + self.__create_base2_frame(column_1, 0, ui_state, video_training_enabled=True) self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=True) - self.__create_noise_frame(column_1, 2, ui_state, callbacks) + self.__create_noise_frame(column_1, 2, ui_state) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_hi_dream_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_hi_dream_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_n_frame(column_0, 1, ui_state, i=1, supports_include=True) self.__create_text_encoder_n_frame(column_0, 2, ui_state, i=2, supports_include=True) self.__create_text_encoder_n_frame(column_0, 3, ui_state, i=3, supports_include=True) self.__create_text_encoder_n_frame(column_0, 4, ui_state, i=4, supports_include=True, supports_layer_skip=False) self.__create_embedding_frame(column_0, 5, ui_state) - self.__create_base2_frame(column_1, 0, ui_state, callbacks, video_training_enabled=True) + self.__create_base2_frame(column_1, 0, ui_state, video_training_enabled=True) self.__create_transformer_frame(column_1, 1, ui_state) - self.__create_noise_frame(column_1, 2, ui_state, callbacks) + self.__create_noise_frame(column_1, 2, ui_state) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __create_base_frame(self, master, row, controller, ui_state, callbacks): + def __create_base_frame(self, master, row, controller, ui_state): frame = self.components.section_frame(master, row) # optimizer self.components.label(frame, 0, 0, "Optimizer", tooltip="The type of optimizer") self.components.options_adv(frame, 0, 1, [str(x) for x in list(Optimizer)], ui_state, "optimizer.optimizer", - command=callbacks.get('restore_optimizer'), - adv_command=callbacks.get('open_optimizer_params')) + command=self.restore_optimizer_config, + adv_command=self.open_optimizer_params) # learning rate scheduler # Wackiness will ensue when reloading configs if we don't check and clear this first. @@ -252,14 +270,12 @@ def __create_base_frame(self, master, row, controller, ui_state, callbacks): tooltip="Learning rate scheduler that automatically changes the learning rate during training") _, d = self.components.options_adv(frame, 1, 1, [str(x) for x in list(LearningRateScheduler)], ui_state, "learning_rate_scheduler", - command=callbacks.get('restore_scheduler'), - adv_command=callbacks.get('open_scheduler_params')) + command=self.restore_scheduler, + adv_command=self.open_scheduler_params) self.lr_scheduler_comp = d['component'] self.lr_scheduler_adv_comp = d['button_component'] # Initial call requires the presence of self.lr_scheduler_adv_comp. - restore_scheduler = callbacks.get('restore_scheduler') - if restore_scheduler: - restore_scheduler(ui_state.get_var("learning_rate_scheduler").get()) + self.restore_scheduler(ui_state.get_var("learning_rate_scheduler").get()) # learning rate self.components.label(frame, 2, 0, "Learning Rate", @@ -308,7 +324,7 @@ def __create_base_frame(self, master, row, controller, ui_state, callbacks): tooltip="Clips the gradient norm. Leave empty to disable gradient clipping.") self.components.entry(frame, 10, 1, ui_state, "clip_grad_norm") - def __create_base2_frame(self, master, row, ui_state, callbacks, video_training_enabled: bool = False, + def __create_base2_frame(self, master, row, ui_state, video_training_enabled: bool = False, supports_circular_padding: bool = False): frame = self.components.section_frame(master, row) row = 0 @@ -338,7 +354,7 @@ def __create_base2_frame(self, master, row, ui_state, callbacks, video_training_ tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") self.components.options_adv(frame, row, 1, [str(x) for x in list(GradientCheckpointingMethod)], ui_state, "gradient_checkpointing", - adv_command=callbacks.get('open_offloading')) + adv_command=self.open_offloading) row += 1 # gradient checkpointing layer offloading @@ -588,7 +604,7 @@ def __create_transformer_frame(self, master, row, ui_state, supports_guidance_sc tooltip="The guidance scale of guidance distilled models passed to the transformer during training.") self.components.entry(frame, 4, 1, ui_state, "transformer.guidance_scale") - def __create_noise_frame(self, master, row, ui_state, callbacks, + def __create_noise_frame(self, master, row, ui_state, supports_generalized_offset_noise: bool = False, supports_dynamic_timestep_shifting: bool = False): frame = self.components.section_frame(master, row) @@ -616,7 +632,7 @@ def __create_noise_frame(self, master, row, ui_state, callbacks, wide_tooltip=True) self.components.options_adv(frame, 3, 1, [str(x) for x in list(TimestepDistribution)], ui_state, "timestep_distribution", - adv_command=callbacks.get('open_timestep_distribution')) + adv_command=self.open_timestep_distribution) # min noising strength self.components.label(frame, 4, 0, "Min Noising Strength", diff --git a/modules/ui/CtkCloudTabView.py b/modules/ui/CtkCloudTabView.py index 0a5249069..ffe37451f 100644 --- a/modules/ui/CtkCloudTabView.py +++ b/modules/ui/CtkCloudTabView.py @@ -9,9 +9,8 @@ class CtkCloudTabView(BaseCloudTabView): def __init__(self, master, controller: CloudTabController, ui_state): - BaseCloudTabView.__init__(self, ctk_components) + BaseCloudTabView.__init__(self, ctk_components, controller) self.master = master - self.controller = controller self.ui_state = ui_state self.frame = ctk.CTkScrollableFrame(master, fg_color="transparent") diff --git a/modules/ui/CtkConceptTabView.py b/modules/ui/CtkConceptTabView.py index 5b3e86ac9..84d8e4508 100644 --- a/modules/ui/CtkConceptTabView.py +++ b/modules/ui/CtkConceptTabView.py @@ -86,6 +86,12 @@ def _maybe_reposition_toolbar(self, width): else: self._toolbar.grid_configure(row=0, column=4, columnspan=2, sticky="ew", padx=10) + def _update_filters(self): + self._create_element_list(search=self.search_var.get(), + type=self.filter_var.get(), + show_disabled=self.show_disabled_var.get()) + self._refresh_show_disabled_text() + def _reset_filters(self): self.search_var.set("") self.filter_var.set("ALL") @@ -109,9 +115,7 @@ class CtkConceptWidgetView(BaseConceptWidgetView, ctk.CTkFrame): def __init__(self, master, concept, i, open_command, remove_command, clone_command, save_command, controller): ctk.CTkFrame.__init__(self, master=master, width=150, height=170, corner_radius=10, bg_color="transparent") - BaseConceptWidgetView.__init__(self, ctk_components) - - self.concept = concept + BaseConceptWidgetView.__init__(self, ctk_components, concept) self.ui_state = CtkUIState(self, concept) self.image_ui_state = CtkUIState(self, concept.image) self.text_ui_state = CtkUIState(self, concept.text) diff --git a/modules/ui/CtkTrainUIView.py b/modules/ui/CtkTrainUIView.py index d3c1a70b5..41af55230 100644 --- a/modules/ui/CtkTrainUIView.py +++ b/modules/ui/CtkTrainUIView.py @@ -82,7 +82,13 @@ class CtkTrainUIView(BaseTrainUIView, ctk.CTk): def __init__(self): ctk.CTk.__init__(self) - BaseTrainUIView.__init__(self, ctk_components) + + train_config = TrainConfig.default_values() + ui_state = CtkUIState(self, train_config) + controller = TrainUIController(train_config) + + BaseTrainUIView.__init__(self, ctk_components, controller, ui_state) + self.controller.view = self self.title("OneTrainer") self.geometry("1100x740") @@ -93,12 +99,6 @@ def __init__(self): ctk.set_appearance_mode("Light" if AppearanceModeTracker.detect_appearance_mode() == 0 else "Dark") ctk.set_default_color_theme("blue") - self.train_config = TrainConfig.default_values() - self.ui_state = CtkUIState(self, self.train_config) - - self.controller = TrainUIController(self.train_config) - self.controller.view = self - self.grid_rowconfigure(0, weight=0) self.grid_rowconfigure(1, weight=1) self.grid_rowconfigure(2, weight=0) @@ -212,7 +212,7 @@ def _set_icon(self): def top_bar(self, master): return CtkTopBarView( master, - TopBarController(self.train_config), + TopBarController(self.controller.train_config), self.ui_state, self.change_model_type, self.change_training_method, @@ -259,7 +259,7 @@ def content_frame(self, master): 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.change_training_method(self.controller.train_config.training_method) return frame @@ -274,7 +274,7 @@ def create_general_tab(self, master): return frame def create_model_tab(self, master): - return CtkModelTabView(master, ModelTabController(self.train_config), self.ui_state) + return CtkModelTabView(master, ModelTabController(self.controller.train_config), self.ui_state) def create_data_tab(self, master): frame = ctk.CTkScrollableFrame(master, fg_color="transparent") @@ -288,13 +288,13 @@ def create_data_tab(self, master): return frame def create_concepts_tab(self, master): - return CtkConceptTabView(master, ConceptTabController(self.train_config), self.ui_state) + return CtkConceptTabView(master, ConceptTabController(self.controller.train_config), self.ui_state) def create_training_tab(self, master) -> CtkTrainingTabView: - return CtkTrainingTabView(master, TrainingTabController(self.train_config), self.ui_state) + return CtkTrainingTabView(master, TrainingTabController(self.controller.train_config), self.ui_state) def create_cloud_tab(self, master) -> CtkCloudTabView: - return CtkCloudTabView(master, CloudTabController(self.train_config, parent=self), self.ui_state) + return CtkCloudTabView(master, CloudTabController(self.controller.train_config, parent=self), self.ui_state) def create_sampling_tab(self, master): master.grid_rowconfigure(0, weight=0) @@ -311,7 +311,7 @@ def create_sampling_tab(self, master): frame = ctk.CTkFrame(master=master, corner_radius=0) frame.grid(row=1, column=0, sticky="nsew") - return CtkSamplingTabView(frame, SamplingTabController(self.train_config), self.ui_state) + return CtkSamplingTabView(frame, SamplingTabController(self.controller.train_config), self.ui_state) def create_backup_tab(self, master): frame = ctk.CTkScrollableFrame(master, fg_color="transparent") @@ -336,7 +336,7 @@ def embedding_tab(self, master): return frame def create_additional_embeddings_tab(self, master): - return CtkAdditionalEmbeddingsTabView(master, AdditionalEmbeddingsTabController(self.train_config), self.ui_state) + return CtkAdditionalEmbeddingsTabView(master, AdditionalEmbeddingsTabController(self.controller.train_config), self.ui_state) def create_tools_tab(self, master): frame = ctk.CTkScrollableFrame(master, fg_color="transparent") @@ -376,7 +376,7 @@ def change_training_method(self, training_method: TrainingMethod): self.tabview.delete("embedding") if training_method == TrainingMethod.LORA and "LoRA" not in self.tabview._tab_dict: - self.lora_tab = CtkLoraTabView(self.tabview.add("LoRA"), LoraTabController(self.train_config), self.ui_state) + self.lora_tab = CtkLoraTabView(self.tabview.add("LoRA"), LoraTabController(self.controller.train_config), self.ui_state) if training_method == TrainingMethod.EMBEDDING and "embedding" not in self.tabview._tab_dict: self.embedding_tab(self.tabview.add("embedding")) diff --git a/modules/ui/CtkTrainingTabView.py b/modules/ui/CtkTrainingTabView.py index bc29488dd..bce309876 100644 --- a/modules/ui/CtkTrainingTabView.py +++ b/modules/ui/CtkTrainingTabView.py @@ -47,31 +47,25 @@ def refresh_ui(self): column_2.grid(row=0, column=2, sticky="nsew") column_2.grid_columnconfigure(0, weight=1) - callbacks = { - 'restore_optimizer': lambda *args: self.controller.restore_optimizer_config(self.ui_state), - 'open_optimizer_params': self._open_optimizer_params_window, - 'restore_scheduler': self._restore_scheduler_config, - 'open_scheduler_params': self._open_scheduler_params_window, - 'open_offloading': self._open_offloading_window, - 'open_timestep_distribution': self._open_timestep_distribution_window, - } - - self.build(column_0, column_1, column_2, self.controller, self.ui_state, callbacks) - - def _restore_scheduler_config(self, variable): + self.build(column_0, column_1, column_2, self.controller, self.ui_state) + + def restore_optimizer_config(self, variable: str): + self.controller.restore_optimizer_config(self.ui_state) + + def open_optimizer_params(self): + self.master.wait_window(self.controller.open_optimizer_params_window(self.master, self.ui_state, CtkOptimizerParamsWindowView)) + + def restore_scheduler(self, variable: str): if not hasattr(self, 'lr_scheduler_adv_comp'): return state = "normal" if self.controller.is_custom_scheduler_value(variable) else "disabled" self.lr_scheduler_adv_comp.configure(state=state) - def _open_optimizer_params_window(self): - self.master.wait_window(self.controller.open_optimizer_params_window(self.master, self.ui_state, CtkOptimizerParamsWindowView)) - - def _open_scheduler_params_window(self): + def open_scheduler_params(self): self.master.wait_window(self.controller.open_scheduler_params_window(self.master, self.ui_state, CtkSchedulerParamsWindowView)) - def _open_timestep_distribution_window(self): - self.master.wait_window(self.controller.open_timestep_distribution_window(self.master, self.ui_state, CtkTimestepDistributionWindowView)) - - def _open_offloading_window(self): + def open_offloading(self): self.master.wait_window(self.controller.open_offloading_window(self.master, self.ui_state, CtkOffloadingWindowView)) + + def open_timestep_distribution(self): + self.master.wait_window(self.controller.open_timestep_distribution_window(self.master, self.ui_state, CtkTimestepDistributionWindowView)) From d03987122553521e733cfd858ffd5eeac5c4d867 Mon Sep 17 00:00:00 2001 From: dxqb Date: Fri, 15 May 2026 11:51:00 +0200 Subject: [PATCH 05/21] refactor: move Windows DPI awareness block to CtkTrainUIView The ctypes DPI awareness call is toolkit-specific (fixes CTK transparency on Windows monitor changes). It already exists in CtkTrainUIView.py and has no place in the toolkit-agnostic controller. Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/TrainUIController.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/modules/ui/TrainUIController.py b/modules/ui/TrainUIController.py index 15d1eaff6..77c4de1a7 100644 --- a/modules/ui/TrainUIController.py +++ b/modules/ui/TrainUIController.py @@ -1,15 +1,12 @@ -import ctypes import datetime import json import os -import platform import subprocess import sys import threading import time import traceback import webbrowser -from contextlib import suppress from pathlib import Path import scripts.generate_debug_report @@ -28,13 +25,6 @@ import torch -# 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": - with suppress(Exception): - # 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 TrainUIController: def __init__(self, config: TrainConfig): From 1bb5a567104ae53ad19bb00594c1922d95510b1e Mon Sep 17 00:00:00 2001 From: dxqb Date: Fri, 15 May 2026 12:36:28 +0200 Subject: [PATCH 06/21] refactor: remove unused __init__ from BaseConceptTabView search_var/filter_var/show_disabled_var were stored in the base but never used there after _update_filters() was removed. Subclasses manage them. Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/BaseConceptTabView.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/modules/ui/BaseConceptTabView.py b/modules/ui/BaseConceptTabView.py index 4d2fca570..6a55e515b 100644 --- a/modules/ui/BaseConceptTabView.py +++ b/modules/ui/BaseConceptTabView.py @@ -15,11 +15,6 @@ class BaseConceptTabView(BaseConfigListView): _FILTER_TYPES = ["ALL", "STANDARD", "VALIDATION", "PRIOR_PREDICTION"] - def __init__(self, search_var, filter_var, show_disabled_var): - self.search_var = search_var - self.filter_var = filter_var - self.show_disabled_var = show_disabled_var - def _element_matches_filters(self, element): if not self.filters.get("show_disabled", True): if hasattr(element, 'enabled') and not element.enabled: From 8dd86c62ce2001739ef37a45dba9d77d8fec1a3c Mon Sep 17 00:00:00 2001 From: dxqb Date: Sat, 16 May 2026 16:18:54 +0200 Subject: [PATCH 07/21] fix: apply upstream COFT removal to CtkLoraTabView Mirrors the change from upstream commit e928ddab (Remove COFT #1447), which removed COFT from LoraTab.py. The merge didn't carry it across the rename to CtkLoraTabView.py. Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/CtkLoraTabView.py | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/modules/ui/CtkLoraTabView.py b/modules/ui/CtkLoraTabView.py index 1c73d90ce..faabf9607 100644 --- a/modules/ui/CtkLoraTabView.py +++ b/modules/ui/CtkLoraTabView.py @@ -121,19 +121,10 @@ def setup_lora(self, peft_type: PeftType): tooltip=f"The block size parameter used when creating a new {name}") components.entry(master, 1, 1, self.ui_state, "oft_block_size", required=True) - # COFT - components.label(master, 1, 3, "Constrained OFT (COFT)", - tooltip="Use the constrained variant of OFT. This constrains the learned rotation to stay very close to the identity matrix, limiting adaptation to only small changes. This improves training stability, helps prevent overfitting on small datasets, and better preserves the base models original knowledge but it may lack expressiveness for tasks requiring substantial adaptation and introduces an additional hyperparameter (COFT Epsilon) that needs tuning.") - components.switch(master, 1, 4, self.ui_state, "oft_coft") - - components.label(master, 2, 3, "COFT Epsilon", - tooltip="The control strength of COFT. Only has an effect if COFT is enabled.") - components.entry(master, 2, 4, self.ui_state, "coft_eps") - # Block Share - components.label(master, 3, 3, "Block Share", + components.label(master, 1, 3, "Block Share", tooltip="Share the OFT parameters between blocks. A single rotation matrix is shared across all blocks within a layer, drastically cutting the number of trainable parameters and yielding very compact adapter files, potentially improving generalization but at the cost of significant expressiveness, which can lead to underfitting on more complex or diverse tasks.") - components.switch(master, 3, 4, self.ui_state, "oft_block_share") + components.switch(master, 1, 4, self.ui_state, "oft_block_share") # Dropout Percentage components.label(master, 2, 0, "Dropout Probability", From c60245812c52e8be686dbf779e62ac5db870f901 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Fri, 29 May 2026 19:41:48 +0200 Subject: [PATCH 08/21] Sync Ctk* copies with upstream changes merged into Base* files Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/CtkLoraTabView.py | 70 ++++++++++++++++++++++++++++++++ modules/ui/CtkTrainUIView.py | 22 ++++++++-- modules/ui/CtkTrainingTabView.py | 26 +++++++----- modules/ui/TrainUIController.py | 22 ++++++++-- 4 files changed, 124 insertions(+), 16 deletions(-) diff --git a/modules/ui/CtkLoraTabView.py b/modules/ui/CtkLoraTabView.py index faabf9607..6e9975476 100644 --- a/modules/ui/CtkLoraTabView.py +++ b/modules/ui/CtkLoraTabView.py @@ -39,6 +39,7 @@ def refresh_ui(self): ("LoRA", PeftType.LORA), ("LoHa", PeftType.LOHA), ("OFT v2", PeftType.OFT_2), + ("LoKr", PeftType.LOKR), ], self.ui_state, "peft_type", command=self.setup_lora) def setup_lora(self, peft_type: PeftType): @@ -46,6 +47,8 @@ def setup_lora(self, peft_type: PeftType): name = "LoHa" elif peft_type == PeftType.OFT_2: name = "OFT v2" + elif peft_type == PeftType.LOKR: + name = "LoKr" else: name = "LoRA" @@ -126,6 +129,11 @@ def setup_lora(self, peft_type: PeftType): tooltip="Share the OFT parameters between blocks. A single rotation matrix is shared across all blocks within a layer, drastically cutting the number of trainable parameters and yielding very compact adapter files, potentially improving generalization but at the cost of significant expressiveness, which can lead to underfitting on more complex or diverse tasks.") components.switch(master, 1, 4, self.ui_state, "oft_block_share") + # Scaled OFT (SOFT) + components.label(master, 2, 3, "Scaled OFT (SOFT)", + tooltip="Applies a scaling factor to the learned weights. This ensures that the effective learning rate remains consistent across different block sizes. Without this, different block sizes require significantly different learning rates.") + components.switch(master, 2, 4, self.ui_state, "oft_scaled") + # Dropout Percentage components.label(master, 2, 0, "Dropout Probability", tooltip="Dropout probability. This percentage of the rotated adapter nodes that will be randomly restored to the base model initial statue. Helps with overfitting. 0 disables, 1 maximum.") @@ -143,3 +151,65 @@ def setup_lora(self, peft_type: PeftType): components.label(master, 4, 0, "Bundle Embeddings", tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") components.switch(master, 4, 1, self.ui_state, "bundle_additional_embeddings") + + # LoKr + elif peft_type == PeftType.LOKR: + # LoKr Main Settings + components.label(master, 1, 0, f"{name} dimension", + tooltip="The dimension parameter used for the secondary decomposition. Analogous to rank in LoRA.") + components.entry(master, 1, 1, self.ui_state, "lokr_dim") + + components.label(master, 2, 0, "Decomposition Factor", + tooltip="Factor for Kronecker product decomposition. -1 for auto, which is recommended. Changing this drastically affects parameter count.") + components.entry(master, 2, 1, self.ui_state, "lokr_decompose_factor") + + # alpha + components.label(master, 3, 0, f"{name} alpha", + tooltip=f"The alpha parameter used when creating a new {name}") + components.entry(master, 3, 1, self.ui_state, "lora_alpha") + + # Dropout Percentage + components.label(master, 4, 0, "Dropout Probability", + tooltip="Dropout probability. This percentage of model nodes will be randomly ignored at each training step. Helps with overfitting. 0 disables, 1 maximum.") + components.entry(master, 4, 1, self.ui_state, "dropout_probability") + + # LoKr weight dtype + components.label(master, 5, 0, f"{name} Weight Data Type", + tooltip=f"The {name} weight data type used for training. This can reduce memory consumption, but reduces precision") + components.options_kv(master, 5, 1, [ + ("float32", DataType.FLOAT_32), + ("bfloat16", DataType.BFLOAT_16), + ], self.ui_state, "lora_weight_dtype") + + # LoKr Vectorization trick + components.label(master, 6, 0, "Kronecker-Vec Trick", + tooltip="Uses an accelerated path that bypasses the materialization of the full Kronecker product. This delivers a massive speedup to the LoKr without sacrificing precision. Highly recommended.") + components.switch(master, 6, 1, self.ui_state, "lokr_vec_trick") + + #LoKr Decomposition Settings + components.label(master, 1, 3, "Decompose Both Matrices", + tooltip="Perform rank decomposition on both Kronecker product matrices (W1 and W2). Only effective for very small dimensions.") + components.switch(master, 1, 4, self.ui_state, "lokr_decompose_both") + + components.label(master, 2, 3, "Use Tucker Decomposition (Conv)", + tooltip="Use Tucker decomposition for convolutional layers. Can be more efficient for some architectures.") + components.switch(master, 2, 4, self.ui_state, "lokr_use_tucker") + + components.label(master, 3, 3, "Force Full Matrix (W2)", + tooltip="Forces the second Kronecker matrix (W2) to be a full matrix, ignoring the dimension setting. For expert use.") + components.switch(master, 3, 4, self.ui_state, "lokr_full_matrix") + + # LoKr DoRA Settings + components.label(master, 4, 3, "Decompose Weights (DoRA)", + tooltip="Apply weight decomposition (DoRA) on top of the LoKr update.") + components.switch(master, 4, 4, self.ui_state, "lokr_weight_decompose") + + components.label(master, 5, 3, "Apply DoRA on Output Axis", + tooltip="Apply the DoRA weight decomposition on the output axis instead of the input axis.") + components.switch(master, 5, 4, self.ui_state, "lokr_dora_on_output") + + + # Additional embeddings + components.label(master, 6, 3, "Bundle Embeddings", + tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") + components.switch(master, 6, 4, self.ui_state, "bundle_additional_embeddings") diff --git a/modules/ui/CtkTrainUIView.py b/modules/ui/CtkTrainUIView.py index ba90d2e64..91c327694 100644 --- a/modules/ui/CtkTrainUIView.py +++ b/modules/ui/CtkTrainUIView.py @@ -132,6 +132,9 @@ def __init__(self): self.training_callbacks = None self.training_commands = None + self.start_time = None + self.start_total_steps = None + self.always_on_tensorboard_subprocess = None self.current_workspace_dir = self.train_config.workspace_dir self._check_start_always_on_tensorboard() @@ -600,14 +603,21 @@ def open_tensorboard(self): webbrowser.open("http://localhost:" + str(self.train_config.tensorboard_port), new=0, autoraise=False) def _calculate_eta_string(self, train_progress: TrainProgress, max_step: int, max_epoch: int) -> str | None: + assert self.start_time is not None and self.start_total_steps is not None + spent_total = time.monotonic() - self.start_time - steps_done = train_progress.epoch * max_step + train_progress.epoch_step + + # calculate steps done in THIS SESSION only + current_total_steps = train_progress.epoch * max_step + train_progress.epoch_step + steps_done_this_session = current_total_steps - self.start_total_steps + remaining_steps = (max_epoch - train_progress.epoch - 1) * max_step + (max_step - train_progress.epoch_step) - total_eta = spent_total / steps_done * remaining_steps - if train_progress.global_step <= 30: + if steps_done_this_session <= 30: return "Estimating ..." + total_eta = spent_total / steps_done_this_session * remaining_steps + td = datetime.timedelta(seconds=total_eta) days = td.days hours, remainder = divmod(td.seconds, 3600) @@ -632,6 +642,10 @@ def delete_eta_label(self): self.eta_label.configure(text="") def on_update_train_progress(self, train_progress: TrainProgress, max_step: int, max_epoch: int): + # capture session start on first progress update + if self.start_total_steps is None: + self.start_total_steps = train_progress.epoch * max_step + train_progress.epoch_step + self.set_step_progress(train_progress.epoch_step, max_step) self.set_epoch_progress(train_progress.epoch, max_epoch) self.set_eta_label(train_progress, max_step, max_epoch) @@ -715,6 +729,8 @@ def __training_thread_function(self): if self.train_config.cloud.enabled: self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) + # Reset session tracking - actual values captured on first progress callback + self.start_total_steps = None self.start_time = time.monotonic() trainer.train() except Exception: diff --git a/modules/ui/CtkTrainingTabView.py b/modules/ui/CtkTrainingTabView.py index bcca11ae9..c839abf73 100644 --- a/modules/ui/CtkTrainingTabView.py +++ b/modules/ui/CtkTrainingTabView.py @@ -438,34 +438,40 @@ def __create_text_encoder_frame(self, master, row, supports_clip_skip=True, supp frame = ctk.CTkFrame(master=master, corner_radius=5) frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") frame.grid_columnconfigure(0, weight=1) + row = 0 if supports_training: - components.label(frame, 0, 0, "Train Text Encoder", + components.label(frame, row, 0, "Train Text Encoder", tooltip="Enables training the text encoder model") - components.switch(frame, 0, 1, self.ui_state, "text_encoder.train") + components.switch(frame, row, 1, self.ui_state, "text_encoder.train") + row += 1 # dropout - components.label(frame, 1, 0, "Caption Dropout Probability", + components.label(frame, row, 0, "Caption Dropout Probability", tooltip="The Probability for dropping the text encoder conditioning") - components.entry(frame, 1, 1, self.ui_state, "text_encoder.dropout_probability") + components.entry(frame, row, 1, self.ui_state, "text_encoder.dropout_probability") + row += 1 if supports_training: # train text encoder epochs - components.label(frame, 2, 0, "Stop Training After", + components.label(frame, row, 0, "Stop Training After", tooltip="When to stop training the text encoder") - components.time_entry(frame, 2, 1, self.ui_state, "text_encoder.stop_training_after", + components.time_entry(frame, row, 1, self.ui_state, "text_encoder.stop_training_after", "text_encoder.stop_training_after_unit", supports_time_units=False) + row += 1 # text encoder learning rate - components.label(frame, 3, 0, "Text Encoder Learning Rate", + components.label(frame, row, 0, "Text Encoder Learning Rate", tooltip="The learning rate of the text encoder. Overrides the base learning rate") - components.entry(frame, 3, 1, self.ui_state, "text_encoder.learning_rate") + components.entry(frame, row, 1, self.ui_state, "text_encoder.learning_rate") + row += 1 if supports_clip_skip: # text encoder layer skip (clip skip) - components.label(frame, 4, 0, "Clip Skip", + components.label(frame, row, 0, "Clip Skip", tooltip="The number of additional clip layers to skip. 0 = the model default") - components.entry(frame, 4, 1, self.ui_state, "text_encoder_layer_skip") + components.entry(frame, row, 1, self.ui_state, "text_encoder_layer_skip") + row += 1 if supports_sequence_length: # text encoder sequence length diff --git a/modules/ui/TrainUIController.py b/modules/ui/TrainUIController.py index ba90d2e64..91c327694 100644 --- a/modules/ui/TrainUIController.py +++ b/modules/ui/TrainUIController.py @@ -132,6 +132,9 @@ def __init__(self): self.training_callbacks = None self.training_commands = None + self.start_time = None + self.start_total_steps = None + self.always_on_tensorboard_subprocess = None self.current_workspace_dir = self.train_config.workspace_dir self._check_start_always_on_tensorboard() @@ -600,14 +603,21 @@ def open_tensorboard(self): webbrowser.open("http://localhost:" + str(self.train_config.tensorboard_port), new=0, autoraise=False) def _calculate_eta_string(self, train_progress: TrainProgress, max_step: int, max_epoch: int) -> str | None: + assert self.start_time is not None and self.start_total_steps is not None + spent_total = time.monotonic() - self.start_time - steps_done = train_progress.epoch * max_step + train_progress.epoch_step + + # calculate steps done in THIS SESSION only + current_total_steps = train_progress.epoch * max_step + train_progress.epoch_step + steps_done_this_session = current_total_steps - self.start_total_steps + remaining_steps = (max_epoch - train_progress.epoch - 1) * max_step + (max_step - train_progress.epoch_step) - total_eta = spent_total / steps_done * remaining_steps - if train_progress.global_step <= 30: + if steps_done_this_session <= 30: return "Estimating ..." + total_eta = spent_total / steps_done_this_session * remaining_steps + td = datetime.timedelta(seconds=total_eta) days = td.days hours, remainder = divmod(td.seconds, 3600) @@ -632,6 +642,10 @@ def delete_eta_label(self): self.eta_label.configure(text="") def on_update_train_progress(self, train_progress: TrainProgress, max_step: int, max_epoch: int): + # capture session start on first progress update + if self.start_total_steps is None: + self.start_total_steps = train_progress.epoch * max_step + train_progress.epoch_step + self.set_step_progress(train_progress.epoch_step, max_step) self.set_epoch_progress(train_progress.epoch, max_epoch) self.set_eta_label(train_progress, max_step, max_epoch) @@ -715,6 +729,8 @@ def __training_thread_function(self): if self.train_config.cloud.enabled: self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) + # Reset session tracking - actual values captured on first progress callback + self.start_total_steps = None self.start_time = time.monotonic() trainer.train() except Exception: From 01cca79010fc7fe862172f463304b4c688eb43a3 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Fri, 29 May 2026 20:42:53 +0200 Subject: [PATCH 09/21] Move Rolling Backup Count under the Rolling Backup toggle Was sitting on the same row at columns 3-4; now on its own row directly below the toggle for clearer visual grouping. Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/BaseTrainUIView.py | 22 +++++++++++----------- modules/ui/CtkTrainUIView.py | 22 +++++++++++----------- 2 files changed, 22 insertions(+), 22 deletions(-) diff --git a/modules/ui/BaseTrainUIView.py b/modules/ui/BaseTrainUIView.py index 91c327694..b9fa0c04a 100644 --- a/modules/ui/BaseTrainUIView.py +++ b/modules/ui/BaseTrainUIView.py @@ -449,32 +449,32 @@ def create_backup_tab(self, master): components.switch(frame, 1, 1, self.ui_state, "rolling_backup") # rolling backup count - components.label(frame, 1, 3, "Rolling Backup Count", + components.label(frame, 2, 0, "Rolling Backup Count", tooltip="Defines the number of backups to keep if rolling backups are enabled") - components.entry(frame, 1, 4, self.ui_state, "rolling_backup_count") + components.entry(frame, 2, 1, self.ui_state, "rolling_backup_count") # backup before save - components.label(frame, 2, 0, "Backup Before Save", + components.label(frame, 3, 0, "Backup Before Save", tooltip="Create a full backup before saving the final model") - components.switch(frame, 2, 1, self.ui_state, "backup_before_save") + components.switch(frame, 3, 1, self.ui_state, "backup_before_save") # save after - components.label(frame, 3, 0, "Save Every", + components.label(frame, 4, 0, "Save Every", tooltip="The interval used when automatically saving the model during training") - components.time_entry(frame, 3, 1, self.ui_state, "save_every", "save_every_unit") + components.time_entry(frame, 4, 1, self.ui_state, "save_every", "save_every_unit") # save now - components.button(frame, 3, 3, "save now", self.save_now) + components.button(frame, 4, 3, "save now", self.save_now) # skip save - components.label(frame, 4, 0, "Skip First", + components.label(frame, 5, 0, "Skip First", tooltip="Start saving automatically after this interval has elapsed") - components.entry(frame, 4, 1, self.ui_state, "save_skip_first", width=50, sticky="nw") + components.entry(frame, 5, 1, self.ui_state, "save_skip_first", width=50, sticky="nw") # save filename prefix - components.label(frame, 5, 0, "Save Filename Prefix", + components.label(frame, 6, 0, "Save Filename Prefix", tooltip="The prefix for filenames used when saving the model during training") - components.entry(frame, 5, 1, self.ui_state, "save_filename_prefix") + components.entry(frame, 6, 1, self.ui_state, "save_filename_prefix") frame.pack(fill="both", expand=1) return frame diff --git a/modules/ui/CtkTrainUIView.py b/modules/ui/CtkTrainUIView.py index 91c327694..b9fa0c04a 100644 --- a/modules/ui/CtkTrainUIView.py +++ b/modules/ui/CtkTrainUIView.py @@ -449,32 +449,32 @@ def create_backup_tab(self, master): components.switch(frame, 1, 1, self.ui_state, "rolling_backup") # rolling backup count - components.label(frame, 1, 3, "Rolling Backup Count", + components.label(frame, 2, 0, "Rolling Backup Count", tooltip="Defines the number of backups to keep if rolling backups are enabled") - components.entry(frame, 1, 4, self.ui_state, "rolling_backup_count") + components.entry(frame, 2, 1, self.ui_state, "rolling_backup_count") # backup before save - components.label(frame, 2, 0, "Backup Before Save", + components.label(frame, 3, 0, "Backup Before Save", tooltip="Create a full backup before saving the final model") - components.switch(frame, 2, 1, self.ui_state, "backup_before_save") + components.switch(frame, 3, 1, self.ui_state, "backup_before_save") # save after - components.label(frame, 3, 0, "Save Every", + components.label(frame, 4, 0, "Save Every", tooltip="The interval used when automatically saving the model during training") - components.time_entry(frame, 3, 1, self.ui_state, "save_every", "save_every_unit") + components.time_entry(frame, 4, 1, self.ui_state, "save_every", "save_every_unit") # save now - components.button(frame, 3, 3, "save now", self.save_now) + components.button(frame, 4, 3, "save now", self.save_now) # skip save - components.label(frame, 4, 0, "Skip First", + components.label(frame, 5, 0, "Skip First", tooltip="Start saving automatically after this interval has elapsed") - components.entry(frame, 4, 1, self.ui_state, "save_skip_first", width=50, sticky="nw") + components.entry(frame, 5, 1, self.ui_state, "save_skip_first", width=50, sticky="nw") # save filename prefix - components.label(frame, 5, 0, "Save Filename Prefix", + components.label(frame, 6, 0, "Save Filename Prefix", tooltip="The prefix for filenames used when saving the model during training") - components.entry(frame, 5, 1, self.ui_state, "save_filename_prefix") + components.entry(frame, 6, 1, self.ui_state, "save_filename_prefix") frame.pack(fill="both", expand=1) return frame From 647afc27b158c320e0ed9a7d7af04abbb07a36c7 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Fri, 29 May 2026 20:45:06 +0200 Subject: [PATCH 10/21] Move Rolling Backup Count under the Rolling Backup toggle Was sitting on the same row at columns 3-4; now on its own row directly below the toggle for clearer visual grouping. Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/BaseTrainUIView.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/modules/ui/BaseTrainUIView.py b/modules/ui/BaseTrainUIView.py index 9a443a342..ff3acd2b7 100644 --- a/modules/ui/BaseTrainUIView.py +++ b/modules/ui/BaseTrainUIView.py @@ -277,32 +277,32 @@ def build_backup_tab_content(self, frame, controller, ui_state): self.components.switch(frame, 1, 1, ui_state, "rolling_backup") # rolling backup count - self.components.label(frame, 1, 3, "Rolling Backup Count", + self.components.label(frame, 2, 0, "Rolling Backup Count", tooltip="Defines the number of backups to keep if rolling backups are enabled") - self.components.entry(frame, 1, 4, ui_state, "rolling_backup_count") + self.components.entry(frame, 2, 1, ui_state, "rolling_backup_count") # backup before save - self.components.label(frame, 2, 0, "Backup Before Save", + self.components.label(frame, 3, 0, "Backup Before Save", tooltip="Create a full backup before saving the final model") - self.components.switch(frame, 2, 1, ui_state, "backup_before_save") + self.components.switch(frame, 3, 1, ui_state, "backup_before_save") # save after - self.components.label(frame, 3, 0, "Save Every", + self.components.label(frame, 4, 0, "Save Every", tooltip="The interval used when automatically saving the model during training") - self.components.time_entry(frame, 3, 1, ui_state, "save_every", "save_every_unit") + self.components.time_entry(frame, 4, 1, ui_state, "save_every", "save_every_unit") # save now - self.components.button(frame, 3, 3, "save now", self.save_now) + self.components.button(frame, 4, 3, "save now", self.save_now) # skip save - self.components.label(frame, 4, 0, "Skip First", + self.components.label(frame, 5, 0, "Skip First", tooltip="Start saving automatically after this interval has elapsed") - self.components.entry(frame, 4, 1, ui_state, "save_skip_first", width=50, sticky="nw") + self.components.entry(frame, 5, 1, ui_state, "save_skip_first", width=50, sticky="nw") # save filename prefix - self.components.label(frame, 5, 0, "Save Filename Prefix", + self.components.label(frame, 6, 0, "Save Filename Prefix", tooltip="The prefix for filenames used when saving the model during training") - self.components.entry(frame, 5, 1, ui_state, "save_filename_prefix") + self.components.entry(frame, 6, 1, ui_state, "save_filename_prefix") def build_embedding_tab_content(self, frame, controller, ui_state): # embedding model name From a0bd3e44ae7d6d0ff3ffa155102f9e7cf71b3da1 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Thu, 4 Jun 2026 23:02:10 +0200 Subject: [PATCH 11/21] Sync CtkConceptWindowView and ConceptWindowController with Base copy Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/ConceptWindowController.py | 3 ++- modules/ui/CtkConceptWindowView.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/modules/ui/ConceptWindowController.py b/modules/ui/ConceptWindowController.py index f58879d5f..824ff43e4 100644 --- a/modules/ui/ConceptWindowController.py +++ b/modules/ui/ConceptWindowController.py @@ -575,7 +575,8 @@ def get_concept_path(path: str) -> str | None: def __download_dataset(self): try: - huggingface_hub.login(token=self.train_config.secrets.huggingface_token, new_session=False) + if self.train_config.secrets.huggingface_token != "": + huggingface_hub.login(token=self.train_config.secrets.huggingface_token) huggingface_hub.snapshot_download(repo_id=self.concept.path, repo_type="dataset") except Exception: traceback.print_exc() diff --git a/modules/ui/CtkConceptWindowView.py b/modules/ui/CtkConceptWindowView.py index f58879d5f..824ff43e4 100644 --- a/modules/ui/CtkConceptWindowView.py +++ b/modules/ui/CtkConceptWindowView.py @@ -575,7 +575,8 @@ def get_concept_path(path: str) -> str | None: def __download_dataset(self): try: - huggingface_hub.login(token=self.train_config.secrets.huggingface_token, new_session=False) + if self.train_config.secrets.huggingface_token != "": + huggingface_hub.login(token=self.train_config.secrets.huggingface_token) huggingface_hub.snapshot_download(repo_id=self.concept.path, repo_type="dataset") except Exception: traceback.print_exc() From fdea894863ea3c7d45defbc912f9dbd0dc2bcc53 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Sat, 13 Jun 2026 11:01:47 +0200 Subject: [PATCH 12/21] Sync CtkConceptWindowView with BaseConceptWindowView after upstream merge Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/CtkConceptWindowView.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/modules/ui/CtkConceptWindowView.py b/modules/ui/CtkConceptWindowView.py index 824ff43e4..f58879d5f 100644 --- a/modules/ui/CtkConceptWindowView.py +++ b/modules/ui/CtkConceptWindowView.py @@ -575,8 +575,7 @@ def get_concept_path(path: str) -> str | None: def __download_dataset(self): try: - if self.train_config.secrets.huggingface_token != "": - huggingface_hub.login(token=self.train_config.secrets.huggingface_token) + huggingface_hub.login(token=self.train_config.secrets.huggingface_token, new_session=False) huggingface_hub.snapshot_download(repo_id=self.concept.path, repo_type="dataset") except Exception: traceback.print_exc() From 1f73b3a6487a4bc87364b0eb017333b3cb25d172 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Fri, 19 Jun 2026 20:30:00 +0200 Subject: [PATCH 13/21] Sync Ctk* views with Base* after merging upstream --- modules/ui/CtkConceptWindowView.py | 3 ++- modules/ui/CtkTimestepDistributionWindowView.py | 2 +- modules/ui/CtkTrainingTabView.py | 2 +- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/modules/ui/CtkConceptWindowView.py b/modules/ui/CtkConceptWindowView.py index f58879d5f..824ff43e4 100644 --- a/modules/ui/CtkConceptWindowView.py +++ b/modules/ui/CtkConceptWindowView.py @@ -575,7 +575,8 @@ def get_concept_path(path: str) -> str | None: def __download_dataset(self): try: - huggingface_hub.login(token=self.train_config.secrets.huggingface_token, new_session=False) + if self.train_config.secrets.huggingface_token != "": + huggingface_hub.login(token=self.train_config.secrets.huggingface_token) huggingface_hub.snapshot_download(repo_id=self.concept.path, repo_type="dataset") except Exception: traceback.print_exc() diff --git a/modules/ui/CtkTimestepDistributionWindowView.py b/modules/ui/CtkTimestepDistributionWindowView.py index 21e41ce3e..313083568 100644 --- a/modules/ui/CtkTimestepDistributionWindowView.py +++ b/modules/ui/CtkTimestepDistributionWindowView.py @@ -133,7 +133,7 @@ def __content_frame(self, master): # dynamic timestep shifting components.label(frame, 6, 0, "Dynamic Timestep Shifting", - tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Dynamic Timestep Shifting is not shown in the preview. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) + tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Dynamic Timestep Shifting is not shown in the preview. Note: For Z-Image, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) components.switch(frame, 6, 1, self.ui_state, "dynamic_timestep_shifting") diff --git a/modules/ui/CtkTrainingTabView.py b/modules/ui/CtkTrainingTabView.py index c839abf73..f897bb8ce 100644 --- a/modules/ui/CtkTrainingTabView.py +++ b/modules/ui/CtkTrainingTabView.py @@ -700,7 +700,7 @@ def __create_noise_frame(self, master, row, supports_generalized_offset_noise: b if supports_dynamic_timestep_shifting: # dynamic timestep shifting components.label(frame, 9, 0, "Dynamic Timestep Shifting", - tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) + tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Note: For Z-Image, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) components.switch(frame, 9, 1, self.ui_state, "dynamic_timestep_shifting") From 826835ede1577651c6bc3f701bb0f78d30d1f8a0 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Sat, 27 Jun 2026 10:33:15 +0200 Subject: [PATCH 14/21] Fix CaptionUI arrow-key navigation and SVD-frame docstring previous_image/next_image called self.view.switch_image, but switch_image lives on the controller, not the view; arrow keys raised AttributeError. Also convert the _make_svd_frames docstring to a # comment. Co-Authored-By: Claude Opus 4.8 --- modules/ui/BaseModelTabView.py | 3 ++- modules/ui/CaptionUIController.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/modules/ui/BaseModelTabView.py b/modules/ui/BaseModelTabView.py index 88f223e22..061207f4a 100644 --- a/modules/ui/BaseModelTabView.py +++ b/modules/ui/BaseModelTabView.py @@ -15,7 +15,8 @@ def __init__(self, components): @abstractmethod def _make_svd_frames(self, parent, row: int): - """Create and place SVDQuant label+entry subframes; return (label_frame, entry_frame).""" + # Create and place SVDQuant label+entry subframes; return (label_frame, entry_frame). + pass def build_content(self, frame, controller, ui_state): if controller.train_config.model_type.is_stable_diffusion(): # TODO simplify diff --git a/modules/ui/CaptionUIController.py b/modules/ui/CaptionUIController.py index 0495322a3..2669b401c 100644 --- a/modules/ui/CaptionUIController.py +++ b/modules/ui/CaptionUIController.py @@ -96,11 +96,11 @@ def switch_image(self, index): def previous_image(self): if len(self.image_rel_paths) > 0 and (self.current_image_index - 1) >= 0: - self.view.switch_image(self.current_image_index - 1) + self.switch_image(self.current_image_index - 1) def next_image(self): if len(self.image_rel_paths) > 0 and (self.current_image_index + 1) < len(self.image_rel_paths): - self.view.switch_image(self.current_image_index + 1) + self.switch_image(self.current_image_index + 1) def load_directory(self, include_subdirectories: bool = False): self.scan_directory(include_subdirectories) From 3f676d76e8b6f09338e381ce3ab05ffc7b78cc5a Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Sat, 27 Jun 2026 11:08:58 +0200 Subject: [PATCH 15/21] Fix standalone caption/convert/video UI launch scripts after Ctk view split The Base/Ctk/Controller refactor deleted modules/ui/CaptionUI.py, ConvertModelUI.py and VideoToolUI.py, but the three entry scripts still imported them, crashing on launch with ModuleNotFoundError. Wire each script through the new controller + create_window(parent, view_cls) pattern (as TrainUIController already does internally). The new views are CTkToplevel, not CTk roots, so launch them under a hidden CTk root and wait_window on the toplevel to avoid a stray empty "tk" window and to exit cleanly on close. Co-Authored-By: Claude Opus 4.8 --- scripts/caption_ui.py | 14 +++++++++++--- scripts/convert_model_ui.py | 14 +++++++++++--- scripts/video_tool_ui.py | 14 +++++++++++--- 3 files changed, 33 insertions(+), 9 deletions(-) diff --git a/scripts/caption_ui.py b/scripts/caption_ui.py index aa380cd9a..3dce599ef 100644 --- a/scripts/caption_ui.py +++ b/scripts/caption_ui.py @@ -2,15 +2,23 @@ script_imports() -from modules.ui.CaptionUI import CaptionUI +from modules.ui.CaptionUIController import CaptionUIController +from modules.ui.CtkCaptionUIView import CtkCaptionUIView from modules.util.args.CaptionUIArgs import CaptionUIArgs +import customtkinter as ctk + def main(): args = CaptionUIArgs.parse_args() - ui = CaptionUI(None, args.dir, args.include_subdirectories) - ui.mainloop() + # CtkCaptionUIView is a CTkToplevel, so it needs a CTk root. Hide the root and run the + # toplevel as a standalone window, tearing the root down once the window is closed. + root = ctk.CTk() + root.withdraw() + ui = CaptionUIController(args.dir, args.include_subdirectories).create_window(root, CtkCaptionUIView) + root.wait_window(ui) + root.destroy() if __name__ == '__main__': diff --git a/scripts/convert_model_ui.py b/scripts/convert_model_ui.py index 3d024042d..09cec86fc 100644 --- a/scripts/convert_model_ui.py +++ b/scripts/convert_model_ui.py @@ -2,12 +2,20 @@ script_imports() -from modules.ui.ConvertModelUI import ConvertModelUI +from modules.ui.ConvertModelUIController import ConvertModelUIController +from modules.ui.CtkConvertModelUIView import CtkConvertModelUIView + +import customtkinter as ctk def main(): - ui = ConvertModelUI(None) - ui.mainloop() + # CtkConvertModelUIView is a CTkToplevel, so it needs a CTk root. Hide the root and run the + # toplevel as a standalone window, tearing the root down once the window is closed. + root = ctk.CTk() + root.withdraw() + ui = ConvertModelUIController().create_window(root, CtkConvertModelUIView) + root.wait_window(ui) + root.destroy() if __name__ == '__main__': diff --git a/scripts/video_tool_ui.py b/scripts/video_tool_ui.py index 99707506f..990f1a60d 100644 --- a/scripts/video_tool_ui.py +++ b/scripts/video_tool_ui.py @@ -2,12 +2,20 @@ script_imports() -from modules.ui.VideoToolUI import VideoToolUI +from modules.ui.CtkVideoToolUIView import CtkVideoToolUIView +from modules.ui.VideoToolUIController import VideoToolUIController + +import customtkinter as ctk def main(): - ui = VideoToolUI(None) - ui.mainloop() + # CtkVideoToolUIView is a CTkToplevel, so it needs a CTk root. Hide the root and run the + # toplevel as a standalone window, tearing the root down once the window is closed. + root = ctk.CTk() + root.withdraw() + ui = VideoToolUIController().create_window(root, CtkVideoToolUIView) + root.wait_window(ui) + root.destroy() if __name__ == '__main__': From 9706006eb382de8a950aaee3e3ec1fd7406d850a Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Sat, 27 Jun 2026 11:27:14 +0200 Subject: [PATCH 16/21] Fix HF dataset download in concept window after Ctk view split The ctk view refactor moved ConceptWindow.__download_dataset into the new ConceptWindowController.download_dataset from a base predating PR #1524, which silently reverted that PR's fix: the method again called huggingface_hub.login(token=..., new_session=False) with no empty-token guard. new_session was removed from login() in huggingface-hub 1.16, so every call raised TypeError, swallowed by the surrounding except, so snapshot_download never ran and dataset download was fully broken. Re-apply #1524 at the new location: only login when a token is configured, and drop new_session. Co-Authored-By: Claude Opus 4.8 --- modules/ui/ConceptWindowController.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/modules/ui/ConceptWindowController.py b/modules/ui/ConceptWindowController.py index c2c486961..12577891c 100644 --- a/modules/ui/ConceptWindowController.py +++ b/modules/ui/ConceptWindowController.py @@ -55,7 +55,8 @@ def get_concept_path(path: str) -> str | None: def download_dataset(self): try: - huggingface_hub.login(token=self.train_config.secrets.huggingface_token, new_session=False) + if self.train_config.secrets.huggingface_token != "": + huggingface_hub.login(token=self.train_config.secrets.huggingface_token) huggingface_hub.snapshot_download(repo_id=self.concept.path, repo_type="dataset") except Exception: traceback.print_exc() From 7377fd7a26167cd774d6c624c59bc463c91be603 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Sat, 27 Jun 2026 11:37:19 +0200 Subject: [PATCH 17/21] Fix stale UIState type reference in ctk_validation The Base/Ctk refactor renamed UIState to BaseUIState (UIState.py) plus CtkUIState (CtkUIState.py) and updated validation.py, but missed the ctk sibling: ctk_validation.py still imported the now-nonexistent UIState under TYPE_CHECKING. Point it at CtkUIState, the concrete state actually passed to these ctk validators (mirroring validation.py -> BaseUIState). Invisible to ruff and the runtime (used in a string annotation, under TYPE_CHECKING, with future annotations), so only a type checker would flag it. Co-Authored-By: Claude Opus 4.8 --- modules/util/ui/ctk_validation.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/modules/util/ui/ctk_validation.py b/modules/util/ui/ctk_validation.py index 5347ba6d4..b77709af3 100644 --- a/modules/util/ui/ctk_validation.py +++ b/modules/util/ui/ctk_validation.py @@ -18,7 +18,7 @@ ) if TYPE_CHECKING: - from modules.util.ui.UIState import UIState + from modules.util.ui.CtkUIState import CtkUIState import customtkinter as ctk @@ -54,7 +54,7 @@ def __init__( self, component: ctk.CTkEntry, var: tk.Variable, - ui_state: UIState, + ui_state: CtkUIState, var_name: str, max_undo: int = DEFAULT_MAX_UNDO, extra_validate: Callable[[str], str | None] | None = None, @@ -258,7 +258,7 @@ def __init__( self, component: ctk.CTkEntry, var: tk.Variable, - ui_state: UIState, + ui_state: CtkUIState, var_name: str, io_type: PathIOType = PathIOType.INPUT, max_undo: int = DEFAULT_MAX_UNDO, From 8f7089789638b04df1736f2e7196bf8910680a2f Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Sat, 27 Jun 2026 11:40:32 +0200 Subject: [PATCH 18/21] Drop customtkinter dependency from toolkit-agnostic Base views BaseConfigListView and BaseAdditionalEmbeddingsTabView imported customtkinter solely for a `-> ctk.CTkToplevel` return annotation on the abstract open_element_window, the lone toolkit reference in the Base layer (the sibling abstract methods return toolkit objects unannotated). Drop the annotation and the import so the Base layer is genuinely toolkit-free; the concrete Ctk subclasses keep their own CTkToplevel annotation. Co-Authored-By: Claude Opus 4.8 --- modules/ui/BaseAdditionalEmbeddingsTabView.py | 4 +--- modules/ui/BaseConfigListView.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/modules/ui/BaseAdditionalEmbeddingsTabView.py b/modules/ui/BaseAdditionalEmbeddingsTabView.py index 744f0abdc..37321ad04 100644 --- a/modules/ui/BaseAdditionalEmbeddingsTabView.py +++ b/modules/ui/BaseAdditionalEmbeddingsTabView.py @@ -2,8 +2,6 @@ from modules.ui.BaseConfigListView import BaseConfigListView from modules.util import path_util -import customtkinter as ctk - class BaseAdditionalEmbeddingsTabView(BaseConfigListView): @@ -14,7 +12,7 @@ def refresh_ui(self): self.widgets_initialized = False self._create_element_list() - def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: + def open_element_window(self, i, ui_state): pass diff --git a/modules/ui/BaseConfigListView.py b/modules/ui/BaseConfigListView.py index 2165e854b..d0d55de2b 100644 --- a/modules/ui/BaseConfigListView.py +++ b/modules/ui/BaseConfigListView.py @@ -6,8 +6,6 @@ from modules.util import path_util from modules.util.path_util import write_json_atomic -import customtkinter as ctk - class BaseConfigListView(ABC): @@ -19,7 +17,7 @@ def create_widget(self, master, element, i, open_command, remove_command, clone_ pass @abstractmethod - def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: + def open_element_window(self, i, ui_state): pass @abstractmethod From d52b8d2e9aad7a478dc92492f698ed725911fdc3 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Sat, 27 Jun 2026 12:26:28 +0200 Subject: [PATCH 19/21] Re-apply aux-optimizer empty-state default fix (#1444) after Ctk view split The ctk view split copied OptimizerParamsWindow into a controller from a base that predated #1444, reverting `if not current_state:` back to `if current_state is None:`. Since muon_adam_config defaults to {} (not None), the empty initial state fell through to from_dict({}), so the Muon/aux-Adam window opened with bare default_values() instead of MUON_AUX_ADAM_DEFAULTS / ADAMW_ADV. Restore the upstream one-line guard. Co-Authored-By: Claude Opus 4.8 --- modules/ui/OptimizerParamsWindowController.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/ui/OptimizerParamsWindowController.py b/modules/ui/OptimizerParamsWindowController.py index 37f838907..3f6771e76 100644 --- a/modules/ui/OptimizerParamsWindowController.py +++ b/modules/ui/OptimizerParamsWindowController.py @@ -35,7 +35,7 @@ def prepare_muon_adam_config(self) -> tuple['TrainOptimizerConfig', Optimizer]: else: defaults = OPTIMIZER_DEFAULT_PARAMETERS[Optimizer.ADAMW_ADV] - if current_state is None: + if not current_state: adam_config.from_dict(defaults) if current_optimizer != Optimizer.MUON: adam_config.optimizer = Optimizer.ADAMW_ADV From 0f361c9a1082918f82fd556ecd610380cba1ae80 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Mon, 29 Jun 2026 23:49:31 +0200 Subject: [PATCH 20/21] Fix sample-window lifecycle: scope binding and defer torch_gc to close connect_window_closed bound unconditionally, so the callback fired once per descendant widget as the toplevel tore down. Guard on the event widget being the window itself so it fires exactly once. With that in place, open_sampling_tool defers torch_gc to window close via connect_window_closed instead of running it inline at open time, restoring the original 'free GPU memory after the sample window closes' behaviour. Co-Authored-By: Claude Opus 4.8 --- modules/ui/CtkTrainUIView.py | 3 ++- modules/ui/TrainUIController.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/modules/ui/CtkTrainUIView.py b/modules/ui/CtkTrainUIView.py index 41af55230..c7c6a18f6 100644 --- a/modules/ui/CtkTrainUIView.py +++ b/modules/ui/CtkTrainUIView.py @@ -201,7 +201,8 @@ def show_window(self, window): window.focus_set() def connect_window_closed(self, window, callback): - window.bind("", lambda _: callback()) + # fires for the toplevel and every descendant widget; only react to the window itself + window.bind("", lambda e: callback() if e.widget is window else None) # --- CTK layout and frame builders --- diff --git a/modules/ui/TrainUIController.py b/modules/ui/TrainUIController.py index a6bf4f21b..817e1e572 100644 --- a/modules/ui/TrainUIController.py +++ b/modules/ui/TrainUIController.py @@ -163,7 +163,7 @@ def open_sampling_tool(self, parent, view_cls): ) window = view_cls(parent, controller) parent.show_window(window) - torch_gc() + parent.connect_window_closed(window, torch_gc) def open_manual_sample_window(self, parent, view_cls): training_callbacks = self.training_callbacks From 57d53bff7f53ee670aea3e1b48f93d955248e6cf Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Mon, 29 Jun 2026 23:49:31 +0200 Subject: [PATCH 21/21] Remove PLAN.md from .gitignore Co-Authored-By: Claude Opus 4.8 --- .gitignore | 1 - 1 file changed, 1 deletion(-) diff --git a/.gitignore b/.gitignore index 6d224a5a4..0b3e5887d 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,4 @@ # development -PLAN.md # claude's planning file, kept local-only .idea *.bak *.pyc