diff --git a/.gitignore b/.gitignore
index f32bc0dc7..0519fc04d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -38,3 +38,5 @@ train.bat
debug_report.log
config_diff.txt
CLAUDE.md
+
+caption_ui_settings.json
diff --git a/modules/module/BaseImageCaptionModel.py b/modules/module/BaseImageCaptionModel.py
index 2dfcf4a87..315142a68 100644
--- a/modules/module/BaseImageCaptionModel.py
+++ b/modules/module/BaseImageCaptionModel.py
@@ -126,6 +126,7 @@ def caption_images(
mode: str = 'fill',
progress_callback: Callable[[int, int], None] = None,
error_callback: Callable[[str], None] = None,
+ is_cancelled: Callable[[], bool] = None,
):
"""
Captions all samples in a list
@@ -141,11 +142,14 @@ def caption_images(
- add: creates a new caption for all samples, appending if a caption already exists
progress_callback (`Callable[[int, int], None]`): called after every processed image
error_callback (`Callable[[str], None]`): called for every exception
+ is_cancelled (`Callable[[], bool]`): checked before each image; if it returns True the run stops
"""
if progress_callback is not None:
progress_callback(0, len(filenames))
for i, filename in enumerate(tqdm(filenames)):
+ if is_cancelled is not None and is_cancelled():
+ break
try:
self.caption_image(filename, initial_caption, caption_prefix, caption_postfix, mode)
except Exception:
@@ -164,6 +168,7 @@ def caption_folder(
progress_callback: Callable[[int, int], None] = None,
error_callback: Callable[[str], None] = None,
include_subdirectories: bool = False,
+ is_cancelled: Callable[[], bool] = None,
):
"""
Captions all samples in a folder
@@ -191,4 +196,5 @@ def caption_folder(
mode=mode,
progress_callback=progress_callback,
error_callback=error_callback,
+ is_cancelled=is_cancelled,
)
diff --git a/modules/module/LMStudioCaptionModel.py b/modules/module/LMStudioCaptionModel.py
new file mode 100644
index 000000000..ae8cc099e
--- /dev/null
+++ b/modules/module/LMStudioCaptionModel.py
@@ -0,0 +1,84 @@
+"""LM Studio captioning backend.
+
+The request flow (base64 image -> OpenAI-compatible /v1/chat/completions,
+stripping ... blocks from reasoning models) is adapted from
+LM_Studio_Image_Captioner by Kazi Shashwata Rahman (MIT License):
+https://github.com/shashwata2020/LM_Studio_Image_Captioner
+"""
+
+import base64
+import os
+import re
+
+from modules.module.BaseImageCaptionModel import BaseImageCaptionModel, CaptionSample
+
+import requests
+
+_MIME_TYPES = {
+ ".png": "image/png",
+ ".jpg": "image/jpeg",
+ ".jpeg": "image/jpeg",
+ ".webp": "image/webp",
+ ".bmp": "image/bmp",
+ ".gif": "image/gif",
+}
+_THINK_RE = re.compile(r".*?", re.DOTALL)
+
+
+class LMStudioCaptionModel(BaseImageCaptionModel):
+ """Captions images through an LM Studio OpenAI-compatible /v1 endpoint."""
+
+ def __init__(self, server_url: str, system_prompt: str = "", user_prompt: str = ""):
+ self.server_url = (server_url or "").strip().rstrip("/")
+ self.system_prompt = system_prompt or ""
+ self.user_prompt = user_prompt or ""
+ self._model_id: str | None = None
+
+ def _resolve_model_id(self) -> str:
+ if self._model_id is None:
+ response = requests.get(f"{self.server_url}/models", timeout=10)
+ response.raise_for_status()
+ models = response.json().get("data") or []
+ self._model_id = models[0]["id"] if models else "local-model"
+ return self._model_id
+
+ @staticmethod
+ def _encode_image(image_filename: str) -> tuple[str, str]:
+ ext = os.path.splitext(image_filename)[1].lower()
+ mime_type = _MIME_TYPES.get(ext, "image/jpeg")
+ with open(image_filename, "rb") as f:
+ encoded = base64.b64encode(f.read()).decode("utf-8")
+ return encoded, mime_type
+
+ def generate_caption(
+ self,
+ caption_sample: CaptionSample,
+ initial_caption: str = "",
+ caption_prefix: str = "",
+ caption_postfix: str = "",
+ ):
+ model_id = self._resolve_model_id()
+ encoded_image, mime_type = self._encode_image(caption_sample.image_filename)
+
+ messages = []
+ if self.system_prompt.strip():
+ messages.append({"role": "system", "content": self.system_prompt})
+ messages.append({
+ "role": "user",
+ "content": [
+ {"type": "image_url",
+ "image_url": {"url": f"data:{mime_type};base64,{encoded_image}"}},
+ {"type": "text", "text": self.user_prompt},
+ ],
+ })
+
+ response = requests.post(
+ f"{self.server_url}/chat/completions",
+ json={"model": model_id, "messages": messages},
+ timeout=None,
+ )
+ response.raise_for_status()
+ raw_text = response.json()["choices"][0]["message"]["content"]
+
+ generated = _THINK_RE.sub("", raw_text).replace("\n", " ").strip()
+ return (caption_prefix + generated + caption_postfix).strip()
diff --git a/modules/ui/CaptionUIController.py b/modules/ui/CaptionUIController.py
index 2669b401c..1d8ea00da 100644
--- a/modules/ui/CaptionUIController.py
+++ b/modules/ui/CaptionUIController.py
@@ -2,9 +2,11 @@
import subprocess
import traceback
+from modules.module.BaseImageCaptionModel import CaptionSample
from modules.module.Blip2Model import Blip2Model
from modules.module.BlipModel import BlipModel
from modules.module.ClipSegModel import ClipSegModel
+from modules.module.LMStudioCaptionModel import LMStudioCaptionModel
from modules.module.MaskByColor import MaskByColor
from modules.module.RembgHumanModel import RembgHumanModel
from modules.module.RembgModel import RembgModel
@@ -236,6 +238,12 @@ def update_mask_draw_radius(self, delta):
multiplier = 1.0 + (delta * 0.05)
self.mask_draw_radius = max(0.0025, self.mask_draw_radius * multiplier)
+ def set_mask_draw_origin(self, event_x, event_y):
+ # Called at the start of a fresh stroke (mouse press) so the first dab is a
+ # dot at the click point instead of a line from where the last stroke ended.
+ self.mask_draw_x = event_x
+ self.mask_draw_y = event_y
+
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
@@ -347,6 +355,24 @@ def load_captioning_model(self, model):
print("loading WD14_VIT_v2 model, this may take a while")
self.captioning_model = WDModel(default_device, torch.float16)
+ def load_lmstudio_captioning_model(self, server_url, system_prompt, user_prompt):
+ self._release_models()
+ self.captioning_model = LMStudioCaptionModel(server_url, system_prompt, user_prompt)
+
+ def caption_current_image(self, server_url, system_prompt, user_prompt):
+ if (
+ len(self.image_rel_paths) == 0
+ or self.current_image_index < 0
+ or self.current_image_index >= len(self.image_rel_paths)
+ ):
+ return None
+
+ image_name = self.image_rel_paths[self.current_image_index]
+ image_path = os.path.join(self.dir, image_name)
+
+ model = LMStudioCaptionModel(server_url, system_prompt, user_prompt)
+ return model.generate_caption(CaptionSample(image_path))
+
def print_help(self):
print(self.help_text)
diff --git a/modules/ui/GenerateCaptionsWindowController.py b/modules/ui/GenerateCaptionsWindowController.py
index 2b2411b28..c418cdf2f 100644
--- a/modules/ui/GenerateCaptionsWindowController.py
+++ b/modules/ui/GenerateCaptionsWindowController.py
@@ -26,3 +26,23 @@ def create_captions(self, model_name, path, initial_caption, caption_prefix, cap
include_subdirectories=include_subdirectories,
)
self.parent.load_image()
+
+ def create_captions_lmstudio(self, server_url, system_prompt, user_prompt, path, mode_str,
+ include_subdirectories, progress_callback=None, error_callback=None,
+ is_cancelled=None):
+ self.parent.load_lmstudio_captioning_model(server_url, system_prompt, user_prompt)
+
+ mode = {
+ "Replace all captions": "replace",
+ "Create if absent": "fill",
+ "Add as new line": "add",
+ }[mode_str]
+
+ self.parent.captioning_model.caption_folder(
+ sample_dir=path,
+ mode=mode,
+ progress_callback=progress_callback,
+ error_callback=error_callback,
+ include_subdirectories=include_subdirectories,
+ is_cancelled=is_cancelled,
+ )
diff --git a/modules/ui/PySide6CaptionUIView.py b/modules/ui/PySide6CaptionUIView.py
index 620e4faab..5fdde9da0 100644
--- a/modules/ui/PySide6CaptionUIView.py
+++ b/modules/ui/PySide6CaptionUIView.py
@@ -1,12 +1,401 @@
-from PySide6.QtWidgets import QDialog, QLabel, QPushButton, QVBoxLayout
+import platform
+import threading
+from modules.ui.BaseCaptionUIView import BaseCaptionUIView
+from modules.ui.PySide6GenerateCaptionsWindowView import PySide6GenerateCaptionsWindowView
+from modules.ui.PySide6GenerateMasksWindowView import PySide6GenerateMasksWindowView
+from modules.util.caption_ui_settings import load_caption_ui_settings
+from modules.util.ui import pyside6_components
+from modules.util.ui.pyside6_util import QtABCMeta
-class PySide6CaptionUIView(QDialog):
- def __init__(self, parent, controller):
+from PIL import Image
+from PIL.ImageQt import ImageQt
+from PySide6.QtCore import Qt, QTimer
+from PySide6.QtGui import QColor, QKeySequence, QPainter, QPen, QPixmap, QShortcut
+from PySide6.QtWidgets import (
+ QCheckBox,
+ QDialog,
+ QFileDialog,
+ QHBoxLayout,
+ QLabel,
+ QLineEdit,
+ QListWidget,
+ QMessageBox,
+ QPushButton,
+ QVBoxLayout,
+ QWidget,
+)
+
+
+class _MaskCanvas(QLabel):
+ """Displays the current image/mask and forwards mouse edits to the controller."""
+
+ def __init__(self, controller, get_alpha, is_editing_enabled, parent=None):
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)
+ self._controller = controller
+ self._get_alpha = get_alpha
+ self._is_editing_enabled = is_editing_enabled
+ self._pixmap_size = (0, 0)
+ self._cursor_pos = None
+ self._cursor_inside = False
+ self.setAlignment(Qt.AlignCenter)
+ self.setMinimumSize(controller.image_size, controller.image_size)
+ self.setMouseTracking(True)
+ self.setCursor(Qt.CrossCursor)
+ # take keyboard focus on click so the [ and ] brush-size keys reach the
+ # dialog (rather than the caption text box) once the user is painting
+ self.setFocusPolicy(Qt.ClickFocus)
+
+ def set_display_pixmap(self, pixmap: QPixmap):
+ self._pixmap_size = (pixmap.width(), pixmap.height())
+ self.setPixmap(pixmap)
+
+ def _brush_display_radius(self) -> float:
+ # draw_mask sizes the brush as radius * max(mask dimensions); the mask is
+ # uniformly scaled to the display pixmap, so max(pixmap) gives display px
+ pw, ph = self._pixmap_size
+ return self._controller.mask_draw_radius * max(pw, ph)
+
+ def _map_to_image(self, pos):
+ pw, ph = self._pixmap_size
+ if pw == 0 or ph == 0:
+ return None
+ offset_x = (self.width() - pw) / 2
+ offset_y = (self.height() - ph) / 2
+ x = pos.x() - offset_x
+ y = pos.y() - offset_y
+ if x < 0 or y < 0 or x >= pw or y >= ph:
+ return None
+ return x, y
+
+ def enterEvent(self, event):
+ self._cursor_inside = True
+ self.update()
+
+ def leaveEvent(self, event):
+ self._cursor_inside = False
+ self.update()
+
+ def mousePressEvent(self, event):
+ self._cursor_pos = event.position()
+ if self._is_editing_enabled():
+ mapped = self._map_to_image(event.position())
+ if mapped is not None:
+ self._controller.set_mask_draw_origin(mapped[0], mapped[1])
+ self._handle(event)
+
+ def mouseMoveEvent(self, event):
+ self._cursor_pos = event.position()
+ self._handle(event)
+ if self._is_editing_enabled():
+ self.update()
+
+ def _handle(self, event):
+ if not self._is_editing_enabled():
+ return
+ mapped = self._map_to_image(event.position())
+ if mapped is None:
+ return
+ buttons = event.buttons()
+ is_left = bool(buttons & Qt.LeftButton)
+ is_right = bool(buttons & Qt.RightButton)
+ if not is_left and not is_right:
+ return
+ self._controller.handle_edit_mask(mapped[0], mapped[1], is_left, is_right, self._get_alpha())
+
+ def wheelEvent(self, event):
+ if not self._is_editing_enabled():
+ return
+ delta = 1 if event.angleDelta().y() > 0 else -1
+ self._controller.update_mask_draw_radius(delta)
+ self.update()
+
+ def paintEvent(self, event):
+ super().paintEvent(event)
+ if not (self._is_editing_enabled() and self._cursor_inside and self._cursor_pos):
+ return
+ if self._pixmap_size == (0, 0):
+ return
+
+ radius = self._brush_display_radius()
+ cx, cy = self._cursor_pos.x(), self._cursor_pos.y()
+
+ painter = QPainter(self)
+ painter.setRenderHint(QPainter.Antialiasing)
+ painter.setBrush(Qt.NoBrush)
+ # draw a dark ring then a dashed light ring so it stays visible on any image
+ painter.setPen(QPen(QColor(0, 0, 0, 200), 1.5))
+ painter.drawEllipse(self._cursor_pos, radius, radius)
+ pen = QPen(QColor(255, 255, 255, 220), 1.0)
+ pen.setStyle(Qt.DashLine)
+ painter.setPen(pen)
+ painter.drawEllipse(self._cursor_pos, radius, radius)
+ # a small center dot marks the exact brush center
+ painter.setPen(QPen(QColor(255, 255, 255, 220), 1.0))
+ painter.drawPoint(int(cx), int(cy))
+ painter.end()
+
+
+class PySide6CaptionUIView(BaseCaptionUIView, QDialog, metaclass=QtABCMeta):
+ def __init__(self, parent, controller):
+ QDialog.__init__(self, parent)
+ BaseCaptionUIView.__init__(self, pyside6_components)
+
+ self.controller = controller
+ controller.view = self
+
+ self._base_title = "OneTrainer - Dataset Tool"
+ self.setWindowTitle(self._base_title)
+ self.resize(1280, 980)
+
+ root = QVBoxLayout(self)
+ root.setContentsMargins(8, 8, 8, 8)
+ root.setSpacing(8)
+
+ root.addWidget(self._build_top_bar())
+
+ body = QHBoxLayout()
+ body.setSpacing(8)
+ root.addLayout(body, 1)
+
+ self.file_list = QListWidget(self)
+ self.file_list.setFixedWidth(300)
+ self.file_list.currentRowChanged.connect(self._on_row_changed)
+ body.addWidget(self.file_list)
+
+ body.addWidget(self._build_content_column(), 1)
+
+ self._install_shortcuts()
+ self.controller.load_directory()
+
+ # ---- layout builders ------------------------------------------------
+
+ def _build_top_bar(self) -> QWidget:
+ bar = QWidget(self)
+ lo = QHBoxLayout(bar)
+ lo.setContentsMargins(0, 0, 0, 0)
+ lo.setSpacing(6)
+
+ def add(text, slot, tooltip):
+ b = QPushButton(text, bar)
+ b.setToolTip(tooltip)
+ b.clicked.connect(slot)
+ lo.addWidget(b)
+ return b
+
+ add("Open", self.open_directory, "open a new directory")
+ add("Generate Masks", self.open_mask_window, "batch-generate masks")
+ add("Generate Captions", self.open_caption_window, "batch-generate captions with LM Studio")
+ self.caption_image_button = add(
+ "Caption Image", self.caption_current_image,
+ "caption the current image with LM Studio")
+ if platform.system() == "Windows":
+ add("Open in Explorer", self.open_in_explorer, "open the current image in Explorer")
+
+ self.include_subdirs_check = QCheckBox("include subdirectories", bar)
+ self.include_subdirs_check.setChecked(
+ bool(self.controller.config_ui_data.get("include_subdirectories", False)))
+ self.include_subdirs_check.toggled.connect(self._on_include_subdirs_toggled)
+ lo.addWidget(self.include_subdirs_check)
+
+ lo.addStretch(1)
+
+ help_button = QPushButton("Help", bar)
+ help_button.setToolTip(self.controller.help_text + "\n [ / ]: decrease / increase brush size")
+ help_button.clicked.connect(self.controller.print_help)
+ lo.addWidget(help_button)
+ return bar
+
+ def _build_content_column(self) -> QWidget:
+ column = QWidget(self)
+ lo = QVBoxLayout(column)
+ lo.setContentsMargins(0, 0, 0, 0)
+ lo.setSpacing(8)
+
+ # mask controls row
+ controls = QWidget(column)
+ controls_lo = QHBoxLayout(controls)
+ controls_lo.setContentsMargins(0, 0, 0, 0)
+ controls_lo.setSpacing(6)
+
+ draw_button = QPushButton("Draw", controls)
+ draw_button.setToolTip("draw a mask using a brush")
+ draw_button.clicked.connect(self.draw_mask_editing_mode)
+ controls_lo.addWidget(draw_button)
+
+ fill_button = QPushButton("Fill", controls)
+ fill_button.setToolTip("draw a mask using a fill tool")
+ fill_button.clicked.connect(self.fill_mask_editing_mode)
+ controls_lo.addWidget(fill_button)
+
+ self.enable_mask_editing_check = QCheckBox("Enable Mask Editing", controls)
+ controls_lo.addWidget(self.enable_mask_editing_check)
+
+ controls_lo.addSpacing(20)
+ controls_lo.addWidget(QLabel("Brush Alpha", controls))
+ self.mask_alpha_edit = QLineEdit("1.0", controls)
+ self.mask_alpha_edit.setFixedWidth(50)
+ controls_lo.addWidget(self.mask_alpha_edit)
+ controls_lo.addStretch(1)
+ lo.addWidget(controls)
+
+ # image canvas
+ self.canvas = _MaskCanvas(
+ self.controller,
+ self._get_alpha,
+ self.enable_mask_editing_check.isChecked,
+ column,
+ )
+ lo.addWidget(self.canvas, 1, Qt.AlignCenter)
+
+ # prompt entry
+ self.prompt_edit = QLineEdit(column)
+ lo.addWidget(self.prompt_edit)
+ return column
+
+ def _install_shortcuts(self):
+ QShortcut(QKeySequence(Qt.Key_Up), self, self.controller.previous_image)
+ QShortcut(QKeySequence(Qt.Key_Down), self, self.controller.next_image)
+ # window-level so Enter saves regardless of whether the prompt box or the
+ # canvas has focus (covers both the main Return and the keypad Enter keys)
+ QShortcut(QKeySequence(Qt.Key_Return), self, self._save)
+ QShortcut(QKeySequence(Qt.Key_Enter), self, self._save)
+ QShortcut(QKeySequence("Ctrl+M"), self, self._toggle_mask)
+ QShortcut(QKeySequence("Ctrl+D"), self, self.draw_mask_editing_mode)
+ QShortcut(QKeySequence("Ctrl+F"), self, self.fill_mask_editing_mode)
+
+ # ---- helpers --------------------------------------------------------
+
+ def _get_alpha(self) -> float:
+ try:
+ return float(self.mask_alpha_edit.text())
+ except ValueError:
+ return 1.0
+
+ def _on_row_changed(self, index):
+ if index >= 0 and index != self.controller.current_image_index:
+ self.controller.switch_image(index)
+
+ def _on_include_subdirs_toggled(self, checked):
+ self.controller.config_ui_data["include_subdirectories"] = checked
+
+ def _save(self):
+ has_image = 0 <= self.controller.current_image_index < len(self.controller.image_rel_paths)
+ self.controller.save(self.prompt_edit.text())
+ if has_image:
+ self.setWindowTitle(f"{self._base_title} — saved ✓")
+ QTimer.singleShot(1500, self, lambda: self.setWindowTitle(self._base_title))
+
+ def _toggle_mask(self):
+ self.controller.toggle_mask()
+ self.refresh_image()
+
+ def _change_brush(self, delta):
+ # positive delta grows the brush, negative shrinks it (same as wheel)
+ self.controller.update_mask_draw_radius(delta)
+ self.canvas.update()
+
+ def keyPressEvent(self, event):
+ # handled here (not as a QShortcut) so typing '[' or ']' into the caption
+ # box still works - a focused text field consumes the key before it reaches
+ # the dialog, and only reaches us when focus is on the canvas/dialog
+ if event.key() == Qt.Key_BracketRight:
+ self._change_brush(1)
+ event.accept()
+ return
+ if event.key() == Qt.Key_BracketLeft:
+ self._change_brush(-1)
+ event.accept()
+ return
+ super().keyPressEvent(event)
+
+ # ---- view callbacks invoked by the controller -----------------------
+
+ def refresh_file_list(self):
+ self.file_list.blockSignals(True)
+ self.file_list.clear()
+ self.file_list.addItems(self.controller.image_rel_paths)
+ self.file_list.blockSignals(False)
+
+ def focus_prompt(self):
+ self.prompt_edit.setFocus()
+
+ def on_image_switched(self, old_index, new_index, prompt):
+ self.file_list.blockSignals(True)
+ self.file_list.setCurrentRow(new_index)
+ self.file_list.blockSignals(False)
+ self.refresh_image()
+ self.prompt_edit.setText(prompt)
+
+ def on_image_cleared(self):
+ blank = Image.new("RGB", (self.controller.image_size, self.controller.image_size), (0, 0, 0))
+ self.canvas.set_display_pixmap(self._to_pixmap(blank))
+
+ def refresh_image(self):
+ pil_image, _size = self.controller.get_display_image()
+ self.canvas.set_display_pixmap(self._to_pixmap(pil_image))
+
+ @staticmethod
+ def _to_pixmap(pil_image) -> QPixmap:
+ return QPixmap.fromImage(ImageQt(pil_image.convert("RGBA")))
+
+ # ---- abstract method implementations --------------------------------
+
+ def open_directory(self):
+ directory = QFileDialog.getExistingDirectory(self, "Select folder", self.controller.dir or "")
+ if directory:
+ self.controller.dir = directory
+ self.controller.load_directory(
+ include_subdirectories=self.controller.config_ui_data["include_subdirectories"])
+
+ def open_mask_window(self):
+ window = self.controller.open_mask_window(self, PySide6GenerateMasksWindowView)
+ window.exec()
+ self.controller.switch_image(self.controller.current_image_index)
+
+ def open_caption_window(self):
+ window = self.controller.open_caption_window(self, PySide6GenerateCaptionsWindowView)
+ window.exec()
+ self.controller.switch_image(self.controller.current_image_index)
+
+ def open_in_explorer(self):
+ self.controller.open_in_explorer()
+
+ def draw_mask_editing_mode(self, *args):
+ self.controller.set_mask_editing_mode("draw")
+
+ def fill_mask_editing_mode(self, *args):
+ self.controller.set_mask_editing_mode("fill")
+
+ # ---- single-image LM Studio captioning ------------------------------
+
+ def caption_current_image(self):
+ settings = load_caption_ui_settings()
+ self.caption_image_button.setEnabled(False)
+
+ def worker():
+ try:
+ caption = self.controller.caption_current_image(
+ settings["server_url"], settings["system_prompt"], settings["user_prompt"])
+ except Exception as e:
+ message = str(e)
+ QTimer.singleShot(0, self, lambda: self._caption_failed(message))
+ return
+ QTimer.singleShot(0, self, lambda: self._caption_done(caption))
+
+ threading.Thread(target=worker, daemon=True).start()
+
+ def _caption_done(self, caption):
+ self.caption_image_button.setEnabled(True)
+ if caption:
+ self.prompt_edit.setText(caption)
+
+ def _caption_failed(self, message):
+ self.caption_image_button.setEnabled(True)
+ QMessageBox.warning(self, "Caption failed", message)
+
+ # ---- lifecycle ------------------------------------------------------
+
+ def closeEvent(self, event):
+ self.controller._release_models()
+ super().closeEvent(event)
diff --git a/modules/ui/PySide6GenerateCaptionsWindowView.py b/modules/ui/PySide6GenerateCaptionsWindowView.py
index 09d82f74b..ffa3e5ff8 100644
--- a/modules/ui/PySide6GenerateCaptionsWindowView.py
+++ b/modules/ui/PySide6GenerateCaptionsWindowView.py
@@ -1,118 +1,188 @@
-import contextlib
-import tkinter as tk
-from tkinter import filedialog
+import threading
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 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 = ""
-
+from modules.util.caption_ui_settings import load_caption_ui_settings, save_caption_ui_settings
+from modules.util.ui.pyside6_util import QtABCMeta
+
+from PySide6.QtCore import Qt, QTimer
+from PySide6.QtWidgets import (
+ QCheckBox,
+ QComboBox,
+ QDialog,
+ QFileDialog,
+ QGridLayout,
+ QLabel,
+ QLineEdit,
+ QProgressBar,
+ QPushButton,
+ QTextEdit,
+ QWidget,
+)
+
+
+class PySide6GenerateCaptionsWindowView(BaseGenerateCaptionsWindowView, QDialog, metaclass=QtABCMeta):
+ def __init__(self, parent, controller: GenerateCaptionsWindowController, path, parent_include_subdirectories):
+ QDialog.__init__(self, parent)
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")
- 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._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)
-
- 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()
+ self._running = False
+ self._cancel_requested = False
+
+ settings = load_caption_ui_settings()
+
+ self.setWindowTitle("Batch generate captions (LM Studio)")
+ self.resize(540, 640)
+
+ modes = ["Replace all captions", "Create if absent", "Add as new line"]
+
+ lo = QGridLayout(self)
+ lo.setContentsMargins(10, 10, 10, 10)
+ lo.setColumnStretch(1, 1)
+ row = 0
+
+ self.url_edit = QLineEdit(settings["server_url"], self)
+ self._add_row(lo, row, "Server URL", self.url_edit)
+ row += 1
+
+ self.path_edit = QLineEdit(path or "", self)
+ self._add_row(lo, row, "Folder", self._path_row(self.path_edit))
+ row += 1
+
+ lo.addWidget(QLabel("System Prompt", self), row, 0, Qt.AlignLeft | Qt.AlignTop)
+ self.system_edit = QTextEdit(self)
+ self.system_edit.setPlainText(settings["system_prompt"])
+ self.system_edit.setFixedHeight(90)
+ lo.addWidget(self.system_edit, row, 1)
+ row += 1
+
+ lo.addWidget(QLabel("User Prompt", self), row, 0, Qt.AlignLeft | Qt.AlignTop)
+ self.user_edit = QTextEdit(self)
+ self.user_edit.setPlainText(settings["user_prompt"])
+ self.user_edit.setFixedHeight(90)
+ lo.addWidget(self.user_edit, row, 1)
+ row += 1
+
+ self.mode_combo = QComboBox(self)
+ self.mode_combo.addItems(modes)
+ self.mode_combo.setCurrentText("Create if absent")
+ self._add_row(lo, row, "Mode", self.mode_combo)
+ row += 1
+
+ self.include_subdirs_check = QCheckBox("Include subfolders", self)
+ self.include_subdirs_check.setChecked(bool(parent_include_subdirectories))
+ lo.addWidget(self.include_subdirs_check, row, 1)
+ row += 1
+
+ lo.addWidget(QLabel("Log", self), row, 0, Qt.AlignLeft | Qt.AlignTop)
+ self.status_box = QTextEdit(self)
+ self.status_box.setReadOnly(True)
+ self.status_box.setFixedHeight(120)
+ lo.addWidget(self.status_box, row, 1)
+ row += 1
+
+ self.progress_label = QLabel("Progress: 0/0", self)
+ self.progress_bar = QProgressBar(self)
+ self.progress_bar.setRange(0, 100)
+ lo.addWidget(self.progress_label, row, 0)
+ lo.addWidget(self.progress_bar, row, 1)
+ row += 1
+
+ self.create_button = QPushButton("Create Captions", self)
+ self.create_button.clicked.connect(self._on_create)
+ lo.addWidget(self.create_button, row, 0)
+
+ self.cancel_button = QPushButton("Cancel", self)
+ self.cancel_button.setEnabled(False)
+ self.cancel_button.clicked.connect(self._on_cancel)
+ lo.addWidget(self.cancel_button, row, 1)
+ row += 1
+
+ lo.setRowStretch(row, 1)
+
+ def _add_row(self, lo, row, label_text, widget):
+ lo.addWidget(QLabel(label_text, self), row, 0, Qt.AlignLeft | Qt.AlignVCenter)
+ lo.addWidget(widget, row, 1)
+
+ def _path_row(self, edit):
+ frame = QWidget(self)
+ frame_lo = QGridLayout(frame)
+ frame_lo.setContentsMargins(0, 0, 0, 0)
+ frame_lo.setColumnStretch(0, 1)
+ frame_lo.addWidget(edit, 0, 0)
+ browse = QPushButton("...", frame)
+ browse.setFixedWidth(40)
+ browse.clicked.connect(self._browse)
+ frame_lo.addWidget(browse, 0, 1)
+ return frame
+
+ def _browse(self):
+ directory = QFileDialog.getExistingDirectory(self, "Select folder", self.path_edit.text())
+ if directory:
+ self.path_edit.setText(directory)
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 _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_str=self.mode_var.get(),
- include_subdirectories=self.include_subdirectories_var.get(),
- )
-
- def destroy(self):
- with contextlib.suppress(tk.TclError):
- self.grab_release()
-
- super().destroy()
+ QTimer.singleShot(0, self, lambda: self._apply_progress(value, max_value))
+
+ def _apply_progress(self, value, max_value):
+ max_value = max(1, max_value)
+ self.progress_bar.setValue(int(value / max_value * 100))
+ self.progress_label.setText(f"Progress: {value}/{max_value}")
+
+ def _log(self, text):
+ QTimer.singleShot(0, self, lambda: self.status_box.append(text))
+
+ def _on_create(self):
+ if self._running:
+ return
+
+ server_url = self.url_edit.text().strip()
+ system_prompt = self.system_edit.toPlainText()
+ user_prompt = self.user_edit.toPlainText()
+ path = self.path_edit.text()
+ mode_str = self.mode_combo.currentText()
+ include_subdirectories = self.include_subdirs_check.isChecked()
+
+ save_caption_ui_settings(server_url, system_prompt, user_prompt)
+
+ self._running = True
+ self._cancel_requested = False
+ self.create_button.setEnabled(False)
+ self.cancel_button.setEnabled(True)
+ self.status_box.clear()
+ self._log(f"Connecting to {server_url} ...")
+
+ def error_callback(filename):
+ self._log(f"Error captioning {filename}")
+
+ def worker():
+ try:
+ self.controller.create_captions_lmstudio(
+ server_url=server_url,
+ system_prompt=system_prompt,
+ user_prompt=user_prompt,
+ path=path,
+ mode_str=mode_str,
+ include_subdirectories=include_subdirectories,
+ progress_callback=self.set_progress,
+ error_callback=error_callback,
+ is_cancelled=lambda: self._cancel_requested,
+ )
+ self._log("Cancelled." if self._cancel_requested else "Done.")
+ except Exception as e:
+ message = str(e)
+ self._log(f"Failed: {message}")
+ finally:
+ QTimer.singleShot(0, self, self._on_done)
+
+ threading.Thread(target=worker, daemon=True).start()
+
+ def _on_cancel(self):
+ if self._running and not self._cancel_requested:
+ self._cancel_requested = True
+ self.cancel_button.setEnabled(False)
+ self._log("Cancelling after the current image ...")
+
+ def _on_done(self):
+ self._running = False
+ self.create_button.setEnabled(True)
+ self.cancel_button.setEnabled(False)
diff --git a/modules/ui/PySide6GenerateMasksWindowView.py b/modules/ui/PySide6GenerateMasksWindowView.py
index 631179fac..9151b9505 100644
--- a/modules/ui/PySide6GenerateMasksWindowView.py
+++ b/modules/ui/PySide6GenerateMasksWindowView.py
@@ -1,141 +1,154 @@
-import contextlib
-import tkinter as tk
-from tkinter import filedialog
+import threading
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 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
-
- 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
- """
- ctk.CTkToplevel.__init__(self, parent, *args, **kwargs)
-
+from modules.util.ui.pyside6_util import QtABCMeta
+
+from PySide6.QtCore import Qt, QTimer
+from PySide6.QtWidgets import (
+ QCheckBox,
+ QComboBox,
+ QDialog,
+ QFileDialog,
+ QGridLayout,
+ QLabel,
+ QLineEdit,
+ QProgressBar,
+ QPushButton,
+ QWidget,
+)
+
+
+class PySide6GenerateMasksWindowView(BaseGenerateMasksWindowView, QDialog, metaclass=QtABCMeta):
+ def __init__(self, parent, controller: GenerateMasksWindowController, path, parent_include_subdirectories):
+ QDialog.__init__(self, parent)
self.controller = controller
- 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._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)
-
- 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()
+ self._running = False
+
+ self.setWindowTitle("Batch generate masks")
+ self.resize(400, 470)
+
+ models = ["ClipSeg", "Rembg", "Rembg-Human", "Hex Color"]
+ modes = ["Replace all masks", "Create if absent", "Add to existing",
+ "Subtract from existing", "Blend with existing"]
+
+ lo = QGridLayout(self)
+ lo.setContentsMargins(10, 10, 10, 10)
+ lo.setColumnStretch(1, 1)
+ row = 0
+
+ self.model_combo = QComboBox(self)
+ self.model_combo.addItems(models)
+ self._add_row(lo, row, "Model", self.model_combo)
+ row += 1
+
+ self.path_edit = QLineEdit(path or "", self)
+ self._add_row(lo, row, "Folder", self._path_row(self.path_edit))
+ row += 1
+
+ self.prompt_edit = QLineEdit(self)
+ self._add_row(lo, row, "Prompt", self.prompt_edit)
+ row += 1
+
+ self.mode_combo = QComboBox(self)
+ self.mode_combo.addItems(modes)
+ self.mode_combo.setCurrentText("Create if absent")
+ self._add_row(lo, row, "Mode", self.mode_combo)
+ row += 1
+
+ self.threshold_edit = QLineEdit("0.3", self)
+ self._add_row(lo, row, "Threshold", self.threshold_edit)
+ row += 1
+
+ self.smooth_edit = QLineEdit("5", self)
+ self._add_row(lo, row, "Smooth", self.smooth_edit)
+ row += 1
+
+ self.expand_edit = QLineEdit("10", self)
+ self._add_row(lo, row, "Expand", self.expand_edit)
+ row += 1
+
+ self.alpha_edit = QLineEdit("1", self)
+ self._add_row(lo, row, "Alpha", self.alpha_edit)
+ row += 1
+
+ self.include_subdirs_check = QCheckBox("Include subfolders", self)
+ self.include_subdirs_check.setChecked(bool(parent_include_subdirectories))
+ lo.addWidget(self.include_subdirs_check, row, 1)
+ row += 1
+
+ self.progress_label = QLabel("Progress: 0/0", self)
+ self.progress_bar = QProgressBar(self)
+ self.progress_bar.setRange(0, 100)
+ lo.addWidget(self.progress_label, row, 0)
+ lo.addWidget(self.progress_bar, row, 1)
+ row += 1
+
+ self.create_button = QPushButton("Create Masks", self)
+ self.create_button.clicked.connect(self._on_create)
+ lo.addWidget(self.create_button, row, 0, 1, 2)
+ row += 1
+
+ lo.setRowStretch(row, 1)
+
+ def _add_row(self, lo, row, label_text, widget):
+ lo.addWidget(QLabel(label_text, self), row, 0, Qt.AlignLeft | Qt.AlignVCenter)
+ lo.addWidget(widget, row, 1)
+
+ def _path_row(self, edit):
+ frame = QWidget(self)
+ frame_lo = QGridLayout(frame)
+ frame_lo.setContentsMargins(0, 0, 0, 0)
+ frame_lo.setColumnStretch(0, 1)
+ frame_lo.addWidget(edit, 0, 0)
+ browse = QPushButton("...", frame)
+ browse.setFixedWidth(40)
+ browse.clicked.connect(self._browse)
+ frame_lo.addWidget(browse, 0, 1)
+ return frame
+
+ def _browse(self):
+ directory = QFileDialog.getExistingDirectory(self, "Select folder", self.path_edit.text())
+ if directory:
+ self.path_edit.setText(directory)
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 _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(),
- )
-
- def destroy(self):
- with contextlib.suppress(tk.TclError):
- self.grab_release()
-
- super().destroy()
+ QTimer.singleShot(0, self, lambda: self._apply_progress(value, max_value))
+
+ def _apply_progress(self, value, max_value):
+ max_value = max(1, max_value)
+ self.progress_bar.setValue(int(value / max_value * 100))
+ self.progress_label.setText(f"Progress: {value}/{max_value}")
+
+ def _on_create(self):
+ if self._running:
+ return
+ self._running = True
+ self.create_button.setEnabled(False)
+
+ args = {
+ "model_name": self.model_combo.currentText(),
+ "path": self.path_edit.text(),
+ "prompt": self.prompt_edit.text(),
+ "mode_str": self.mode_combo.currentText(),
+ "alpha_str": self.alpha_edit.text(),
+ "threshold_str": self.threshold_edit.text(),
+ "smooth_str": self.smooth_edit.text(),
+ "expand_str": self.expand_edit.text(),
+ "include_subdirectories": self.include_subdirs_check.isChecked(),
+ }
+
+ def worker():
+ try:
+ self.controller.create_masks(**args)
+ except Exception as e:
+ message = str(e)
+ QTimer.singleShot(0, self, lambda: self.progress_label.setText(f"Error: {message}"))
+ finally:
+ QTimer.singleShot(0, self, self._on_done)
+
+ threading.Thread(target=worker, daemon=True).start()
+
+ def _on_done(self):
+ self._running = False
+ self.create_button.setEnabled(True)
diff --git a/modules/util/caption_ui_settings.py b/modules/util/caption_ui_settings.py
new file mode 100644
index 000000000..01f6ce1ad
--- /dev/null
+++ b/modules/util/caption_ui_settings.py
@@ -0,0 +1,37 @@
+import json
+from contextlib import suppress
+
+from modules.util import path_util
+
+CAPTION_UI_SETTINGS_FILE = path_util.canonical_join(".", "caption_ui_settings.json")
+
+DEFAULT_SERVER_URL = "http://localhost:1234/v1"
+DEFAULT_SYSTEM_PROMPT = (
+ "You write concise, factual captions for training images. "
+ "Reply with a single caption and no extra commentary."
+)
+DEFAULT_USER_PROMPT = "Describe this image in one concise caption."
+
+
+def load_caption_ui_settings() -> dict:
+ settings = {
+ "server_url": DEFAULT_SERVER_URL,
+ "system_prompt": DEFAULT_SYSTEM_PROMPT,
+ "user_prompt": DEFAULT_USER_PROMPT,
+ }
+ with suppress(Exception), open(CAPTION_UI_SETTINGS_FILE, encoding="utf-8") as f:
+ data = json.load(f)
+ if isinstance(data, dict):
+ for key in settings:
+ if isinstance(data.get(key), str):
+ settings[key] = data[key]
+ return settings
+
+
+def save_caption_ui_settings(server_url: str, system_prompt: str, user_prompt: str) -> None:
+ with suppress(Exception):
+ path_util.write_json_atomic(CAPTION_UI_SETTINGS_FILE, {
+ "server_url": server_url,
+ "system_prompt": system_prompt,
+ "user_prompt": user_prompt,
+ })