implement Qt6 views

.gitignore

@@ -38,4 +38,4 @@ pixi.toml
 train.bat
 debug_report.log
 config_diff.txt
-PLAN.md
+CLAUDE.md

modules/ui/BaseConfigListView.py

@@ -149,7 +149,7 @@ def _toggle_items(self):
 
     def __create_configs_dropdown(self):
         if self.configs_dropdown is not None:
-            self.configs_dropdown.destroy()
+            self._destroy_widget(self.configs_dropdown)
 
         self.configs_dropdown = self.components.options_kv(
             self.top_frame, 0, 0, self.configs, self.ui_state, self.attr_name, self.__load_current_config

modules/ui/BaseTopBarView.py

@@ -30,7 +30,7 @@ def _setup_frame_column_weight(self):
         pass
 
     @abstractmethod
-    def _forget_dropdown(self):
+    def _forget_dropdown(self, widget):
         pass
 
     @abstractmethod
@@ -80,11 +80,11 @@ def build(
         # self.components.icon_button(self.frame, 0, 2, "-", self.__remove_config)
 
         # Wiki button
-        self.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, sticky="vew")
 
         # save button
         self.components.button(self.frame, 0, 3, "Save config", self.__save_config,
-                               tooltip="Save the current configuration in a custom preset", width=90)
+                               tooltip="Save the current configuration in a custom preset", width=90, sticky="vew")
 
         # padding
         self._setup_frame_column_weight()
@@ -98,11 +98,12 @@ def build(
             ui_state=ui_state,
             var_name="model_type",
             command=self.__change_model_type,
+            sticky="vew",
         )
 
     def __create_training_method(self):
         if self.training_method:
-            self.training_method.destroy()
+            self._forget_dropdown(self.training_method)
 
         values = self.controller.get_training_methods(self.controller.train_config.model_type)
 
@@ -114,6 +115,7 @@ def __create_training_method(self):
             ui_state=self.ui_state,
             var_name="training_method",
             command=self.change_training_method_callback,
+            sticky="vew",
         )
 
     def __change_model_type(self, model_type: ModelType):
@@ -122,10 +124,11 @@ def __change_model_type(self, model_type: ModelType):
 
     def __create_configs_dropdown(self):
         if self.configs_dropdown is not None:
-            self._forget_dropdown()
+            self._forget_dropdown(self.configs_dropdown)
 
         self.configs_dropdown = self.components.options_kv(
-            self.frame, 0, 1, self.configs, self.config_ui_state, "config_name", self.__load_current_config
+            self.frame, 0, 1, self.configs, self.config_ui_state, "config_name", self.__load_current_config,
+            sticky="vew",
         )
 
     def __save_config(self):

modules/ui/CtkTopBarView.py

@@ -36,8 +36,8 @@ def _get_dropdown_text(self, widget) -> str:
     def _setup_frame_column_weight(self):
         self.frame.grid_columnconfigure(5, weight=1)
 
-    def _forget_dropdown(self):
-        self.configs_dropdown.grid_forget()
+    def _forget_dropdown(self, widget):
+        widget.destroy()
 
     def _show_save_dialog(self, default_value: str, callback):
         dialogs.StringInputDialog(

modules/ui/PySide6AdditionalEmbeddingsTabView.py

@@ -1,17 +1,16 @@
-
 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
+from modules.ui.PySide6ConfigListView import PySide6ConfigListView
+from modules.util.ui import pyside6_components
+from modules.util.ui.PySide6UIState import PySide6UIState
 
-import customtkinter as ctk
+from PySide6.QtWidgets import QWidget
 
 
-class CtkAdditionalEmbeddingsTabView(CtkConfigListView, BaseAdditionalEmbeddingsTabView):
+class PySide6AdditionalEmbeddingsTabView(PySide6ConfigListView, BaseAdditionalEmbeddingsTabView):
 
     def __init__(self, master, controller: AdditionalEmbeddingsTabController, ui_state):
-        CtkConfigListView.__init__(
+        PySide6ConfigListView.__init__(
             self, master, controller, ui_state,
             attr_name="additional_embeddings",
             enable_key="train",
@@ -22,30 +21,34 @@ def __init__(self, master, controller: AdditionalEmbeddingsTabController, ui_sta
         )
 
     def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command):
-        return CtkEmbeddingWidgetView(master, element, i, open_command, remove_command, clone_command, save_command, self.controller)
+        return PySide6EmbeddingWidgetView(master, element, i, open_command, remove_command, clone_command, save_command, self.controller)
 
 
-class CtkEmbeddingWidgetView(BaseEmbeddingWidgetView, ctk.CTkFrame):
+class PySide6EmbeddingWidgetView(BaseEmbeddingWidgetView, QWidget):
 
     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)
+        QWidget.__init__(self, master)
+        BaseEmbeddingWidgetView.__init__(self, pyside6_components)
 
         self.element = element
-        ui_state = CtkUIState(self, element)
+        ui_state = PySide6UIState(element)
 
-        self.grid_columnconfigure(0, weight=1)
+        pyside6_components._layout(self).setColumnStretch(0, 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)
+        top_frame = QWidget(self)
+        pyside6_components._layout(top_frame).setColumnStretch(3, 1)
+        pyside6_components._layout(top_frame).setColumnStretch(5, 1)
+        pyside6_components._layout(self).addWidget(top_frame, 0, 0)
 
-        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)
+        bottom_frame = QWidget(self)
+        pyside6_components._layout(bottom_frame).setColumnStretch(7, 1)
+        pyside6_components._layout(self).addWidget(bottom_frame, 1, 0)
 
         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")
+        pyside6_components._layout(self.parent()).addWidget(self, getattr(self, 'visible_index', self.i), 0)
+        self.show()
+
+    def destroy(self):
+        self.deleteLater()

modules/ui/PySide6CaptionUIView.py

@@ -1,228 +1,12 @@
-from tkinter import filedialog
-
-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
-
-import customtkinter as ctk
-from customtkinter import ScalingTracker, ThemeManager
-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)
-
-        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")
-        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.controller.load_directory()
-
-        self.wait_visibility()
-        self.focus_set()
-        self.after(200, lambda: set_window_icon(self))
-
-    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.controller.image_rel_paths):
-            def __create_switch_image(index):
-                def __switch_image(event):
-                    self.controller.switch_image(index)
-
-                return __switch_image
-
-            label = ctk.CTkLabel(self.file_list, text=filename)
-            label.bind("<Button-1>", __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)
-
-        self.build_mask_buttons(right_frame)
-
-        # 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.controller.image_size, self.controller.image_size)
-        )
-        self.image_label = ctk.CTkLabel(
-            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")
-
-        self.image_label.bind("<Motion>", self.edit_mask)
-        self.image_label.bind("<Button-1>", self.edit_mask)
-        self.image_label.bind("<Button-3>", 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("<Down>", lambda e: self.controller.next_image())
-        component.bind("<Up>", lambda e: self.controller.previous_image())
-        component.bind("<Return>", self.save)
-        component.bind("<Control-m>", self.toggle_mask)
-        component.bind("<Control-d>", self.draw_mask_editing_mode)
-        component.bind("<Control-f>", self.fill_mask_editing_mode)
-
-    def refresh_file_list(self):
-        self.file_list_column(self.bottom_frame)
-
-    def focus_prompt(self):
-        self.prompt_component.focus_set()
-
-    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)
-
-    def on_image_cleared(self):
-        image = Image.new("RGB", (512, 512), (0, 0, 0))
-        self.image.configure(light_image=image)
-
-    def refresh_image(self):
-        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):
-        self.controller.update_mask_draw_radius(delta)
-
-    def edit_mask(self, event):
-        if not self.enable_mask_editing_var.get():
-            return
-
-        if event.widget != self.image_label.children["!label"]:
-            return
-
-        display_scaling = ScalingTracker.get_window_scaling(self)
-
-        event_x = event.x / display_scaling
-        event_y = event.y / display_scaling
-
-        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
-
-        try:
-            alpha = float(self.mask_editing_alpha.get())
-        except Exception:
-            alpha = 1.0
-
-        self.controller.handle_edit_mask(event_x, event_y, is_left, is_right, alpha)
-
-    def save(self, event):
-        self.controller.save(self.prompt_var.get())
-
-    def draw_mask_editing_mode(self, *args):
-        self.controller.set_mask_editing_mode('draw')
-
-        if args:
-            # disable default event
-            return "break"
-        return None
-
-    def fill_mask_editing_mode(self, *args):
-        self.controller.set_mask_editing_mode('fill')
-
-    def toggle_mask(self, *args):
-        self.controller.toggle_mask()
-        self.refresh_image()
-
-    def open_directory(self):
-        new_dir = filedialog.askdirectory()
-
-        if new_dir:
-            self.controller.dir = new_dir
-            self.controller.load_directory(include_subdirectories=self.controller.config_ui_data["include_subdirectories"])
-
-    def open_mask_window(self):
-        self.wait_window(self.controller.open_mask_window(self, CtkGenerateMasksWindowView))
-        self.controller.switch_image(self.controller.current_image_index)
-
-    def open_caption_window(self):
-        self.wait_window(self.controller.open_caption_window(self, CtkGenerateCaptionsWindowView))
-        self.controller.switch_image(self.controller.current_image_index)
-
-    def open_in_explorer(self):
-        self.controller.open_in_explorer()
-
-    def destroy(self):
-        self.controller._release_models()
-        super().destroy()
+from PySide6.QtWidgets import QDialog, QLabel, QPushButton, QVBoxLayout
+
+
+class PySide6CaptionUIView(QDialog):
+    def __init__(self, parent, controller):
+        super().__init__(parent)
+        self.setWindowTitle("Dataset Tool")
+        lo = QVBoxLayout(self)
+        lo.addWidget(QLabel("The dataset tool has not been ported to Qt6 yet.\nYou can still use it by launching the CustomTkinter UI: scripts/train_ui_ctk.py"))
+        ok = QPushButton("OK")
+        ok.clicked.connect(self.accept)
+        lo.addWidget(ok)