Merge PR Nerogar#1445 into preview

Author: dxqb

.gitignore

@@ -1,4 +1,5 @@
 # development
+PLAN.md  # claude's planning file, kept local-only
 .idea
 *.bak
 *.pyc

modules/ui/AdditionalEmbeddingsTab.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

modules/ui/AdditionalEmbeddingsTabController.py

@@ -0,0 +1,75 @@
+
+from modules.ui.BaseConfigListView import BaseConfigListView
+from modules.util import path_util
+
+import customtkinter as ctk
+
+
+class BaseAdditionalEmbeddingsTabView(BaseConfigListView):
+
+    def refresh_ui(self):
+        if self.element_list is not None:
+            self._destroy_frame(self.element_list)
+            self.element_list = None
+        self.widgets_initialized = False
+        self._create_element_list()
+
+    def open_element_window(self, i, ui_state) -> ctk.CTkToplevel:
+        pass
+
+
+class BaseEmbeddingWidgetView:
+
+    def __init__(self, components):
+        self.components = components
+
+    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
+
+        # close button
+        self.components.colored_icon_button(top_frame, 0, 0, "X", "#C00000", lambda: remove_command(self.i))
+
+        # clone button
+        self.components.colored_icon_button(top_frame, 0, 1, "+", "#00C000", lambda: clone_command(self.i, controller.randomize_uuid), padx=5)
+
+        # embedding model names
+        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=path_util.json_path_modifier
+        )
+
+        # 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
+        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
+        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
+        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
+        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
+        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

modules/ui/BaseAdditionalEmbeddingsTabView.py

@@ -0,0 +1,51 @@
+import platform
+from abc import ABC, abstractmethod
+
+
+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")
+        self.components.button(frame, 0, 1, "Generate Masks", self.open_mask_window,
+                          tooltip="open a dialog to automatically generate masks")
+        self.components.button(frame, 0, 2, "Generate Captions", self.open_caption_window,
+                          tooltip="open a dialog to automatically generate captions")
+
+        if platform.system() == "Windows":
+            self.components.button(frame, 0, 3, "Open in Explorer", self.open_in_explorer,
+                              tooltip="open the current image in Explorer")
+
+        self.components.switch(frame, 0, 4, ui_state, "include_subdirectories",
+                          text="include subdirectories")
+
+        frame.grid_columnconfigure(5, weight=1)
+
+        self.components.button(frame, 0, 6, "Help", controller.print_help,
+                          tooltip=controller.help_text)
+
+    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")
+        self.components.button(right_frame, 0, 1, "Fill", self.fill_mask_editing_mode,
+                          tooltip="draw a mask using a fill tool")