Pyside6 dataset tool including lmstudio

.gitignore

@@ -38,3 +38,5 @@ train.bat
 debug_report.log
 config_diff.txt
 CLAUDE.md
+
+caption_ui_settings.json

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,
         )

modules/module/LMStudioCaptionModel.py

@@ -0,0 +1,84 @@
+"""LM Studio captioning backend.
+
+The request flow (base64 image -> OpenAI-compatible /v1/chat/completions,
+stripping <think>...</think> 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"<think>.*?</think>", 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()

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)
 

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,
+        )