diff --git a/.github/workflows/typos.yaml b/.github/workflows/typos.yaml
index a2d02db38..56018c29f 100644
--- a/.github/workflows/typos.yaml
+++ b/.github/workflows/typos.yaml
@@ -18,4 +18,4 @@ jobs:
       - uses: actions/checkout@v4
 
       - name: typos-action
-        uses: crate-ci/typos@v1.32.0
+        uses: crate-ci/typos@v1.34.0
diff --git a/README.md b/README.md
index 252653aca..7d8182dd4 100644
--- a/README.md
+++ b/README.md
@@ -27,7 +27,8 @@ Support for Linux and macOS is also available. While Linux support is actively m
     - [Colab](#-colab)
     - [Runpod, Novita, Docker](#runpod-novita-docker)
 - [Custom Path Defaults](#custom-path-defaults)
-    - [LoRA](#lora)
+- [Server Options](#server-options)
+- [LoRA](#lora)
   - [Sample image generation during training](#sample-image-generation-during-training)
   - [Troubleshooting](#troubleshooting)
   - [Page File Limit](#page-file-limit)
@@ -101,9 +102,9 @@ I would like to express my gratitude to camenduru for their valuable contributio
 
 These options are for users running training on hosted GPU infrastructure or containers.
 
-- **[Runpod setup](docs/runpod_setup.md)** – Ready-made GPU background training via templates.
-- **[Novita setup](docs/novita_setup.md)** – Similar to Runpod, but integrated into the Novita UI.
-- **[Docker setup](docs/docker.md)** – For developers/sysadmins using containerized environments.
+- **[Runpod setup](docs/installation_runpod.md)** – Ready-made GPU background training via templates.
+- **[Novita setup](docs/installation_novita.md)** – Similar to Runpod, but integrated into the Novita UI.
+- **[Docker setup](docs/installation_docker.md)** – For developers/sysadmins using containerized environments.
 
 
 ## Custom Path Defaults with `config.toml`
@@ -180,6 +181,21 @@ If you prefer to name your configuration file differently or store it in another
 
 By effectively using `config.toml`, you can significantly speed up your training setup process. Always refer to the `config example.toml` for the most up-to-date list of configurable paths.
 
+## Server Options
+
+The `config.toml` file can also be used to configure the Gradio server.
+
+### `allowed_paths`
+
+The `allowed_paths` option allows you to specify a list of directories that the application can access. This is useful if you want to store your models, datasets, or other files on an external drive or in a location outside of the application's root directory.
+
+**Example:**
+
+```toml
+[server]
+allowed_paths = ["/mnt/external_drive/models", "/home/user/datasets"]
+```
+
 ## LoRA
 
 To train a LoRA, you can currently use the `train_network.py` code. You can create a LoRA network by using the all-in-one GUI.
diff --git a/config example.toml b/config example.toml
index 3ab9609a7..cbfc76006 100644
--- a/config example.toml	
+++ b/config example.toml	
@@ -4,6 +4,9 @@
 [settings]
 use_shell = false # Use shell during process run of sd-scripts oython code. Most secure is false but some systems may require it to be true to properly run sd-scripts.
 
+[server]
+allowed_paths = [] # A list of paths that the application is allowed to access.
+
 # Default folders location
 [model]
 models_dir = "./models"                    # Pretrained model name or path
diff --git a/gui-uv.sh b/gui-uv.sh
index c6436d1e5..98e54450c 100755
--- a/gui-uv.sh
+++ b/gui-uv.sh
@@ -52,4 +52,40 @@ if [[ "$uv_quiet" == "--quiet" ]]; then
 fi
 
 git submodule update --init --recursive
+
+# --- bitsandbytes compilation (optional, similar to Colab setup) ---
+# The pyproject.toml specifies bitsandbytes, which uv will install.
+# However, if you encounter issues with the pre-compiled version,
+# you might need to compile it from source, similar to the Colab notebook.
+# Uncomment and adapt the following lines if needed.
+# Ensure you have the necessary build tools (like CUDA toolkit for cuda11x).
+
+# BUILD_BITSANDBYTES_FROM_SOURCE=false # Set to true to enable
+# if [ "$BUILD_BITSANDBYTES_FROM_SOURCE" = true ]; then
+#   echo "Attempting to build bitsandbytes from source..."
+#   if [ ! -d "bitsandbytes" ]; then
+#     git clone -b 0.41.0 https://github.com/TimDettmers/bitsandbytes
+#   fi
+#   cd bitsandbytes || exit 1
+#   # IMPORTANT: Adjust CUDA_VERSION if necessary for your setup (e.g., 118 for CUDA 11.8, 12x for CUDA 12.x)
+#   # The Colab notebook used CUDA_VERSION=118
+#   # Common options: cuda11x (for 11.0-11.8), cuda12x (for 12.0-12.x)
+#   # Ensure the corresponding CUDA toolkit is installed and in PATH.
+#   # For ROCm, the build process is different.
+#   # Check bitsandbytes documentation for the correct make target for your GPU/driver.
+#   # Example for CUDA 11.8:
+#   # CUDA_VERSION=118 make cuda11x
+#   # Example for CUDA 12.1:
+#   # make cuda12x
+#   # Then, install using uv within the environment (or pip if uv is not active yet for this part)
+#   # Assuming uv environment is active or will be activated by the main command:
+#   uv pip install .
+#   # Or, if building before activating the main uv environment:
+#   # python setup.py install
+#   cd .. || exit 1
+#   echo "bitsandbytes build attempt finished."
+# fi
+# --- end of bitsandbytes compilation ---
+
+echo "Launching Kohya GUI via uv..."
 uv run $uv_quiet kohya_gui.py --noverify "${args[@]}"
diff --git a/kohya_gui.py b/kohya_gui.py
index e8d64a884..095491837 100644
--- a/kohya_gui.py
+++ b/kohya_gui.py
@@ -102,6 +102,7 @@ def UI(**kwargs):
         "share": False if kwargs.get("do_not_share", False) else kwargs.get("share", False),
         "root_path": kwargs.get("root_path", None),
         "debug": kwargs.get("debug", False),
+        "allowed_paths": config.allowed_paths,
     }
   
     # This line filters out any key-value pairs from `launch_params` where the value is `None`, ensuring only valid parameters are passed to the `launch` function.
diff --git a/kohya_gui/class_gui_config.py b/kohya_gui/class_gui_config.py
index 33064a9da..92c20dac9 100644
--- a/kohya_gui/class_gui_config.py
+++ b/kohya_gui/class_gui_config.py
@@ -16,6 +16,7 @@ def __init__(self, config_file_path: str = "./config.toml"):
         Initialize the KohyaSSGUIConfig class.
         """
         self.config = self.load_config(config_file_path=config_file_path)
+        self.allowed_paths = self.get_allowed_paths()
 
     def load_config(self, config_file_path: str = "./config.toml") -> dict:
         """
@@ -81,6 +82,12 @@ def get(self, key: str, default=None):
         log.debug(f"Returned {data}")
         return data
 
+    def get_allowed_paths(self) -> list:
+        """
+        Retrieves the list of allowed paths from the [server] section of the config file.
+        """
+        return self.get("server.allowed_paths", [])
+
     def is_config_loaded(self) -> bool:
         """
         Checks if the configuration was loaded from a file.
diff --git a/kohya_gui/class_tensorboard.py b/kohya_gui/class_tensorboard.py
index 001c894da..d82789638 100644
--- a/kohya_gui/class_tensorboard.py
+++ b/kohya_gui/class_tensorboard.py
@@ -3,12 +3,21 @@
 import subprocess
 import time
 import webbrowser
+import shutil
+import cpuinfo
 
-try:
-    os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
-    import tensorflow  # Attempt to import tensorflow to check if it is installed
+def check_avx_support():
+    try:
+        info = cpuinfo.get_cpu_info()
+        return 'avx' in info['flags']
+    except Exception:
+        return False
 
-    visibility = True
+try:
+    if shutil.which("tensorboard") and check_avx_support():
+        visibility = True
+    else:
+        visibility = False
 except ImportError:
     visibility = False
 
diff --git a/kohya_gui/kontext_manual_caption_gui.py b/kohya_gui/kontext_manual_caption_gui.py
new file mode 100644
index 000000000..0be377b52
--- /dev/null
+++ b/kohya_gui/kontext_manual_caption_gui.py
@@ -0,0 +1,506 @@
+import gradio as gr
+from easygui import boolbox
+from .common_gui import get_folder_path, scriptdir, list_dirs
+from math import ceil
+import os
+import re
+
+from .custom_logging import setup_logging
+
+# Set up logging
+log = setup_logging()
+
+IMAGES_TO_SHOW = 5
+IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg", ".webp", ".bmp")
+
+
+def _get_caption_path(image_file, images_dir, caption_ext):
+    """
+    Returns the expected path of a caption file for a given image path
+    """
+    if not image_file:
+        return None
+    # REFACTOR: Use os.path.basename to ensure we only have the filename
+    base_name = os.path.basename(image_file)
+    caption_file_name = os.path.splitext(base_name)[0] + caption_ext
+    caption_file_path = os.path.join(images_dir, caption_file_name)
+    return caption_file_path
+
+
+def _get_quick_tags(quick_tags_text):
+    """
+    Gets a list of tags from the quick tags text box
+    """
+    quick_tags = [t.strip() for t in quick_tags_text.split(",") if t.strip()]
+    quick_tags_set = set(tag.lower() for tag in quick_tags) # REFACTOR: Use lowercase for matching
+    return quick_tags, quick_tags_set
+
+
+def _get_tag_checkbox_updates(caption, quick_tags, quick_tags_set):
+    """
+    Updates a list of caption checkboxes to show possible tags and tags
+    already included in the caption
+    """
+    caption_tags_have = [c.strip() for c in caption.split(",") if c.strip()]
+    # REFACTOR: Match case-insensitively against quick_tags_set
+    caption_tags_unique = [t for t in caption_tags_have if t.lower() not in quick_tags_set]
+    caption_tags_all = quick_tags + caption_tags_unique
+    return gr.CheckboxGroup(choices=caption_tags_all, value=caption_tags_have)
+
+
+def paginate_go(page, max_page):
+    try:
+        page_num = int(page) # REFACTOR: Use int, pages are not fractional
+        return paginate(page_num, max_page, 0)
+    except (ValueError, TypeError):
+        gr.Warning(f"Invalid page number: {page}")
+        return gr.update()
+
+
+def derive_target_folder(control_folder):
+    if not control_folder or not os.path.exists(control_folder):
+        return ""
+
+    parent_dir = os.path.dirname(control_folder)
+
+    for item in os.listdir(parent_dir):
+        if os.path.isdir(os.path.join(parent_dir, item)):
+            if "target" in item.lower():
+                return os.path.join(parent_dir, item)
+
+    return os.path.join(parent_dir, "target")
+
+
+def paginate(page, max_page, page_change):
+    # REFACTOR: Ensure page is treated as an integer
+    return int(max(min(int(page) + page_change, max_page), 1))
+
+
+def save_caption(caption, caption_ext, image_file, images_dir):
+    caption_path = _get_caption_path(image_file, images_dir, caption_ext)
+    if caption_path:
+        # REFACTOR: Use 'w' which is sufficient and standard for writing/overwriting a file.
+        with open(caption_path, "w", encoding="utf-8") as f:
+            f.write(caption)
+        log.info(f"Wrote captions to {caption_path}")
+        return gr.Markdown(f"💾 Caption saved to `{caption_path}`", visible=True)
+    return gr.Markdown(visible=False)
+
+
+def delete_images_and_caption(
+    image_file, control_images_dir, target_images_dir, caption_ext
+):
+    if not image_file:
+        return gr.Markdown(visible=False)
+
+    # Delete control image
+    control_image_path = os.path.join(control_images_dir, image_file)
+    if os.path.exists(control_image_path):
+        os.remove(control_image_path)
+        log.info(f"Deleted control image: {control_image_path}")
+
+    # Delete target image
+    target_image_path = os.path.join(target_images_dir, image_file)
+    if os.path.exists(target_image_path):
+        os.remove(target_image_path)
+        log.info(f"Deleted target image: {target_image_path}")
+
+    # Delete caption file
+    caption_path = _get_caption_path(
+        image_file, target_images_dir, caption_ext
+    )
+    if caption_path and os.path.exists(caption_path):
+        os.remove(caption_path)
+        log.info(f"Deleted caption file: {caption_path}")
+
+    return gr.Markdown(
+        f"🗑️ Deleted files for `{image_file}`", visible=True
+    )
+
+
+def update_quick_tags(quick_tags_text, *image_caption_texts):
+    quick_tags, quick_tags_set = _get_quick_tags(quick_tags_text)
+    return [
+        _get_tag_checkbox_updates(caption, quick_tags, quick_tags_set)
+        for caption in image_caption_texts
+    ]
+
+
+def update_image_caption(
+    quick_tags_text, caption, image_file, images_dir, caption_ext, auto_save_is_checked
+):
+    # REFACTOR: Changed parameter name to avoid shadowing built-in
+    if auto_save_is_checked:
+        save_caption(caption, caption_ext, image_file, images_dir)
+
+    quick_tags, quick_tags_set = _get_quick_tags(quick_tags_text)
+    return _get_tag_checkbox_updates(caption, quick_tags, quick_tags_set)
+
+
+def update_image_tags(
+    quick_tags_text,
+    selected_tags,
+    image_file,
+    images_dir,
+    caption_ext,
+    auto_save_is_checked,
+):
+    # REFACTOR: Changed parameter name
+    # Try to determine order by quick tags
+    quick_tags, quick_tags_set = _get_quick_tags(quick_tags_text)
+    selected_tags_set = set(selected_tags)
+
+    output_tags = [t for t in quick_tags if t in selected_tags_set] + [
+        t for t in selected_tags if t not in quick_tags_set
+    ]
+    caption = ", ".join(output_tags)
+
+    if auto_save_is_checked:
+        save_caption(caption, caption_ext, image_file, images_dir)
+
+    return caption
+
+
+def import_tags_from_captions(
+    images_dir,
+    caption_ext,
+    quick_tags_text,
+    ignore_load_tags_word_count,
+    ask_overwrite=True,
+):
+    if not images_dir or not os.path.exists(images_dir):
+        gr.Warning(
+            "Image folder is not set or does not exist. Please load images first."
+        )
+        return gr.update()
+
+    if not caption_ext:
+        gr.Warning("Please provide an extension for the caption files.")
+        return gr.update()
+
+    if quick_tags_text and ask_overwrite:
+        if not boolbox(
+            "Are you sure you wish to overwrite the current quick tags?",
+            choices=("Yes", "No"),
+        ):
+            return gr.update()
+
+    # REFACTOR: Directly iterate over files from os.scandir for slight performance gain
+    tags = []
+    tags_set = set()
+    for entry in os.scandir(images_dir):
+        if entry.is_file() and entry.name.lower().endswith(IMAGE_EXTENSIONS):
+            caption_file_path = _get_caption_path(entry.name, images_dir, caption_ext)
+            if os.path.exists(caption_file_path):
+                with open(caption_file_path, "r", encoding="utf-8") as f:
+                    caption = f.read()
+                    for tag in caption.split(","):
+                        tag = tag.strip()
+                        if not tag:
+                            continue
+                        tag_key = tag.lower()
+                        if tag_key not in tags_set:
+                            # REFACTOR: Simplified word count logic
+                            if tag.count(" ") + 1 <= ignore_load_tags_word_count:
+                                tags.append(tag)
+                                tags_set.add(tag_key)
+
+    # Ensure tags are alphabetically sorted (case-insensitive)
+    tags_sorted = sorted(tags, key=lambda t: t.lower())
+    gr.Info(f"Imported {len(tags_sorted)} unique tags.")
+    return ", ".join(tags_sorted)
+
+
+def load_images(
+    target_images_dir,
+    control_images_dir,
+    caption_ext,
+    quick_tags_text,
+    ignore_load_tags_word_count,
+):
+    def error_message(msg):
+        gr.Warning(msg)
+        # REFACTOR: Return updates for all outputs to clear state on error
+        return [None, None, None, 1, 1, gr.Markdown(f"⚠️ {msg}", visible=True), gr.update()]
+
+    if not target_images_dir or not os.path.exists(target_images_dir):
+        return error_message("Target image folder is missing or does not exist.")
+    if not control_images_dir or not os.path.exists(control_images_dir):
+        return error_message("Control image folder is missing or does not exist.")
+    if not caption_ext:
+        return error_message("Please provide an extension for the caption files.")
+
+    target_image_files = {
+        f
+        for f in os.listdir(target_images_dir)
+        if f.lower().endswith(IMAGE_EXTENSIONS)
+    }
+    control_image_files = {
+        f
+        for f in os.listdir(control_images_dir)
+        if f.lower().endswith(IMAGE_EXTENSIONS)
+    }
+    # REFACTOR: Sort files here once and store them
+    shared_files = sorted(list(target_image_files.intersection(control_image_files)))
+
+    if not shared_files:
+        return error_message(
+            "No shared images found between the target and control directories."
+        )
+
+    total_images = len(shared_files)
+    max_pages = ceil(total_images / IMAGES_TO_SHOW)
+
+    info = f"✅ Loaded {total_images} shared images. {max_pages} pages total."
+    gr.Info(info)
+
+    # Import tags
+    new_quick_tags = import_tags_from_captions(
+        target_images_dir,
+        caption_ext,
+        quick_tags_text,
+        ignore_load_tags_word_count,
+        ask_overwrite=False,
+    )
+
+    # REFACTOR: Return the computed file list to be stored in gr.State
+    return [
+        shared_files,
+        target_images_dir,
+        control_images_dir,
+        1,
+        max_pages,
+        gr.Markdown(info, visible=True),
+        new_quick_tags,
+    ]
+
+
+def update_images(
+    image_files, # REFACTOR: Receive the list of files from gr.State
+    target_images_dir,
+    control_images_dir,
+    caption_ext,
+    quick_tags_text,
+    page,
+):
+    # REFACTOR: No more os.listdir here! We get the list directly from state.
+    if not image_files or not target_images_dir:
+        # Return empty updates if state is not ready
+        empty_row = gr.Row(visible=False)
+        return [empty_row] * (IMAGES_TO_SHOW * 5 + 2)
+
+    quick_tags, quick_tags_set = _get_quick_tags(quick_tags_text or "")
+
+    outputs = []
+    start_index = (int(page) - 1) * IMAGES_TO_SHOW
+
+    # Build component updates in lists
+    rows_update, files_update, target_paths, control_paths, captions, tags_checks = [], [], [], [], [], []
+
+    for i in range(IMAGES_TO_SHOW):
+        image_index = start_index + i
+        is_visible = image_index < len(image_files)
+
+        rows_update.append(gr.Row(visible=is_visible))
+
+        image_file, target_path, control_path, caption = None, None, None, ""
+        if is_visible:
+            image_file = image_files[image_index]
+            target_path = os.path.join(target_images_dir, image_file)
+            control_path = os.path.join(control_images_dir, image_file)
+            caption_file_path = _get_caption_path(image_file, target_images_dir, caption_ext)
+            if caption_file_path and os.path.exists(caption_file_path):
+                with open(caption_file_path, "r", encoding="utf-8") as f:
+                    caption = f.read()
+
+        files_update.append(image_file)
+        target_paths.append(target_path)
+        control_paths.append(control_path)
+        captions.append(caption)
+        tags_checks.append(_get_tag_checkbox_updates(caption, quick_tags, quick_tags_set))
+
+    # Combine all updates into a single list
+    outputs.extend(rows_update)
+    outputs.extend(files_update)
+    outputs.extend(target_paths)
+    outputs.extend(control_paths)
+    outputs.extend(captions)
+    outputs.extend(tags_checks)
+    outputs.extend([gr.Row(visible=True), gr.Row(visible=True)]) # Pagination rows
+
+    return outputs
+
+
+# Gradio UI
+def gradio_kontext_manual_caption_gui_tab(headless=False, default_images_dir=None):
+    from .common_gui import create_refresh_button
+
+    default_images_dir = default_images_dir or os.path.join(scriptdir, "data")
+
+    # REFACTOR: Simplify directory update logic
+    def update_dir_list(path):
+        # FIX: Convert generator from list_dirs to a list
+        return gr.Dropdown(choices=[""] + list(list_dirs(path)))
+
+    # REFACTOR: Define pagination UI and logic in one place
+    def render_pagination_with_logic(page, max_page):
+        with gr.Row(visible=False) as pagination_row:
+            gr.Button("◀ Prev").click(paginate, inputs=[page, max_page, gr.Number(-1, visible=False)], outputs=[page])
+            page_count = gr.Text("Page 1 / 1", show_label=False, interactive=False, text_align="center")
+            page_goto_text = gr.Textbox(show_label=False, placeholder="Go to page...", container=False, scale=1)
+            gr.Button("Next ▶").click(paginate, inputs=[page, max_page, gr.Number(1, visible=False)], outputs=[page])
+            page_goto_text.submit(paginate_go, inputs=[page_goto_text, max_page], outputs=[page])
+        return pagination_row, page_count
+
+    with gr.Tab("Kontext Manual Captioning"):
+        gr.Markdown("This utility allows quick captioning and tagging of images for fine-tuning with before and after images.")
+
+        # REFACTOR: Use gr.State for non-UI data
+        image_files_state = gr.State([])
+
+        info_box = gr.Markdown(visible=False)
+        page = gr.Number(value=1, visible=False)
+        max_page = gr.Number(value=1, visible=False)
+        loaded_images_dir = gr.Text(visible=False)
+        loaded_control_images_dir = gr.Text(visible=False)
+
+        with gr.Group():
+            with gr.Row():
+                # FIX: Convert generator from list_dirs to a list
+                control_images_dir = gr.Dropdown(label="Control image folder", choices=[""] + list(list_dirs(default_images_dir)), value="", interactive=True, allow_custom_value=True, scale=2)
+                create_refresh_button(control_images_dir, lambda: None, lambda: {"choices": list(list_dirs(control_images_dir.value or default_images_dir))}, "open_folder_small")
+                gr.Button("📂", elem_id="open_folder_small", elem_classes=["tool"], visible=not headless).click(get_folder_path, outputs=control_images_dir, show_progress=False)
+                
+                # FIX: Convert generator from list_dirs to a list
+                target_images_dir = gr.Dropdown(label="Target image folder", choices=[""] + list(list_dirs(default_images_dir)), value="", interactive=True, allow_custom_value=True, scale=2)
+                create_refresh_button(target_images_dir, lambda: None, lambda: {"choices": list(list_dirs(target_images_dir.value or default_images_dir))}, "open_folder_small")
+                gr.Button("📂", elem_id="open_folder_small", elem_classes=["tool"], visible=not headless).click(get_folder_path, outputs=target_images_dir, show_progress=False)
+
+            with gr.Row():
+                caption_ext = gr.Dropdown(label="Caption file extension", choices=[".cap", ".caption", ".txt"], value=".txt", interactive=True, allow_custom_value=True)
+                auto_save = gr.Checkbox(label="Autosave", value=True, interactive=True)
+
+            load_images_button = gr.Button("Load Images", variant="primary")
+
+            target_images_dir.change(update_dir_list, inputs=target_images_dir, outputs=target_images_dir, show_progress=False)
+            control_images_dir.change(
+                lambda path, current_target: (
+                    update_dir_list(path),
+                    derive_target_folder(path)
+                    if not current_target
+                    else current_target,
+                ),
+                inputs=[control_images_dir, target_images_dir],
+                outputs=[control_images_dir, target_images_dir],
+                show_progress=False,
+            )
+
+        with gr.Group():
+            quick_tags_text = gr.Textbox(label="Quick Tags", placeholder="Comma separated list of tags", interactive=True)
+            with gr.Row():
+                ignore_load_tags_word_count = gr.Slider(minimum=1, maximum=100, value=10, step=1, label="Ignore Imported Tags Above Word Count", interactive=True, scale=2)
+            
+            with gr.Row():
+                import_tags_button = gr.Button("Import tags from captions", scale=1)
+
+        pagination_row1, page_count1 = render_pagination_with_logic(page, max_page)
+
+        image_rows, image_files, target_image_images, control_image_images, image_caption_texts, image_tag_checks, save_buttons, delete_buttons = [], [], [], [], [], [], [], []
+        for i in range(IMAGES_TO_SHOW):
+            with gr.Row(visible=False) as row:
+                image_file = gr.Text(visible=False)
+                with gr.Column():
+                    control_image_image = gr.Image(type="filepath", label="Control Image")
+                with gr.Column():
+                    target_image_image = gr.Image(type="filepath", label="Target Image")
+                with gr.Column(scale=2):
+                    image_caption_text = gr.TextArea(label="Captions", placeholder="Input captions for target image", interactive=True)
+                    tag_checkboxes = gr.CheckboxGroup([], label="Tags", interactive=True)
+                with gr.Column(min_width=40):
+                    save_button = gr.Button("💾", elem_id="save_button", visible=False)
+                    delete_button = gr.Button("🗑️", elem_id="delete_button")
+
+                image_rows.append(row); image_files.append(image_file); control_image_images.append(control_image_image); target_image_images.append(target_image_image)
+                image_caption_texts.append(image_caption_text); image_tag_checks.append(tag_checkboxes); save_buttons.append(save_button); delete_buttons.append(delete_button)
+
+                image_caption_text.input(update_image_caption, inputs=[quick_tags_text, image_caption_text, image_file, loaded_images_dir, caption_ext, auto_save], outputs=tag_checkboxes)
+                tag_checkboxes.input(update_image_tags, inputs=[quick_tags_text, tag_checkboxes, image_file, loaded_images_dir, caption_ext, auto_save], outputs=[image_caption_text])
+                save_button.click(save_caption, inputs=[image_caption_text, caption_ext, image_file, loaded_images_dir], outputs=info_box)
+                delete_button.click(
+                    delete_images_and_caption,
+                    inputs=[
+                        image_file,
+                        loaded_control_images_dir,
+                        loaded_images_dir,
+                        caption_ext,
+                    ],
+                    outputs=info_box,
+                ).then(
+                    load_images,
+                    inputs=[
+                        loaded_images_dir,
+                        loaded_control_images_dir,
+                        caption_ext,
+                        quick_tags_text,
+                        ignore_load_tags_word_count,
+                    ],
+                    outputs=[
+                        image_files_state,
+                        loaded_images_dir,
+                        loaded_control_images_dir,
+                        page,
+                        max_page,
+                        info_box,
+                        quick_tags_text,
+                    ],
+                )
+
+        pagination_row2, page_count2 = render_pagination_with_logic(page, max_page)
+
+        quick_tags_text.change(update_quick_tags, inputs=[quick_tags_text] + image_caption_texts, outputs=image_tag_checks)
+        import_tags_button.click(
+            import_tags_from_captions,
+            inputs=[
+                loaded_images_dir,
+                caption_ext,
+                quick_tags_text,
+                ignore_load_tags_word_count,
+            ],
+            outputs=quick_tags_text,
+        )
+
+        load_images_outputs = [
+            image_files_state,
+            loaded_images_dir,
+            loaded_control_images_dir,
+            page,
+            max_page,
+            info_box,
+            quick_tags_text,
+        ]
+        load_images_button.click(
+            load_images,
+            inputs=[
+                target_images_dir,
+                control_images_dir,
+                caption_ext,
+                quick_tags_text,
+                ignore_load_tags_word_count,
+            ],
+            outputs=load_images_outputs,
+        )
+
+        # REFACTOR: A single trigger to update the images view
+        update_trigger_inputs = [image_files_state, loaded_images_dir, loaded_control_images_dir, caption_ext, quick_tags_text, page]
+        update_outputs = (
+            image_rows + image_files + target_image_images + control_image_images +
+            image_caption_texts + image_tag_checks + [pagination_row1, pagination_row2]
+        )
+
+        # Trigger update when page or data changes
+        page.change(update_images, inputs=update_trigger_inputs, outputs=update_outputs, show_progress=False)
+        image_files_state.change(update_images, inputs=update_trigger_inputs, outputs=update_outputs, show_progress=False)
+
+        auto_save.change(lambda is_auto_save: [gr.Button(visible=not is_auto_save)] * IMAGES_TO_SHOW, inputs=auto_save, outputs=save_buttons)
+        page.change(lambda p, m: (f"Page {int(p)} / {int(m)}", f"Page {int(p)} / {int(m)}"), inputs=[page, max_page], outputs=[page_count1, page_count2], show_progress=False)
+        max_page.change(lambda p, m: (f"Page {int(p)} / {int(m)}", f"Page {int(p)} / {int(m)}"), inputs=[page, max_page], outputs=[page_count1, page_count2], show_progress=False)
diff --git a/kohya_gui/manual_caption_gui.py b/kohya_gui/manual_caption_gui.py
index 0a31e084d..e630b0891 100644
--- a/kohya_gui/manual_caption_gui.py
+++ b/kohya_gui/manual_caption_gui.py
@@ -19,7 +19,10 @@ def _get_caption_path(image_file, images_dir, caption_ext):
     """
     Returns the expected path of a caption file for a given image path
     """
-    caption_file_name = os.path.splitext(image_file)[0] + caption_ext
+    if not image_file:
+        return None
+    base_name = os.path.basename(image_file)
+    caption_file_name = os.path.splitext(base_name)[0] + caption_ext
     caption_file_path = os.path.join(images_dir, caption_file_name)
     return caption_file_path
 
@@ -29,7 +32,7 @@ def _get_quick_tags(quick_tags_text):
     Gets a list of tags from the quick tags text box
     """
     quick_tags = [t.strip() for t in quick_tags_text.split(",") if t.strip()]
-    quick_tags_set = set(quick_tags)
+    quick_tags_set = set(tag.lower() for tag in quick_tags)
     return quick_tags, quick_tags_set
 
 
@@ -39,30 +42,32 @@ def _get_tag_checkbox_updates(caption, quick_tags, quick_tags_set):
     already included in the caption
     """
     caption_tags_have = [c.strip() for c in caption.split(",") if c.strip()]
-    caption_tags_unique = [t for t in caption_tags_have if t not in quick_tags_set]
+    caption_tags_unique = [t for t in caption_tags_have if t.lower() not in quick_tags_set]
     caption_tags_all = quick_tags + caption_tags_unique
     return gr.CheckboxGroup(choices=caption_tags_all, value=caption_tags_have)
 
 
 def paginate_go(page, max_page):
     try:
-        page = float(page)
-    except:
-        msgbox(f"Invalid page num: {page}")
-        return
-    return paginate(page, max_page, 0)
+        page_num = int(page)
+        return paginate(page_num, max_page, 0)
+    except (ValueError, TypeError):
+        gr.Warning(f"Invalid page number: {page}")
+        return gr.update()
 
 
 def paginate(page, max_page, page_change):
-    return int(max(min(page + page_change, max_page), 1))
+    return int(max(min(int(page) + page_change, max_page), 1))
 
 
 def save_caption(caption, caption_ext, image_file, images_dir):
     caption_path = _get_caption_path(image_file, images_dir, caption_ext)
-    with open(caption_path, "w+", encoding="utf-8") as f:
-        f.write(caption)
-
-    log.info(f"Wrote captions to {caption_path}")
+    if caption_path:
+        with open(caption_path, "w", encoding="utf-8") as f:
+            f.write(caption)
+        log.info(f"Wrote captions to {caption_path}")
+        return gr.Markdown(f"💾 Caption saved to `{caption_path}`", visible=True)
+    return gr.Markdown(visible=False)
 
 
 def update_quick_tags(quick_tags_text, *image_caption_texts):
@@ -74,9 +79,9 @@ def update_quick_tags(quick_tags_text, *image_caption_texts):
 
 
 def update_image_caption(
-    quick_tags_text, caption, image_file, images_dir, caption_ext, auto_save
+    quick_tags_text, caption, image_file, images_dir, caption_ext, auto_save_is_checked
 ):
-    if auto_save:
+    if auto_save_is_checked:
         save_caption(caption, caption_ext, image_file, images_dir)
 
     quick_tags, quick_tags_set = _get_quick_tags(quick_tags_text)
@@ -89,7 +94,7 @@ def update_image_tags(
     image_file,
     images_dir,
     caption_ext,
-    auto_save,
+    auto_save_is_checked,
 ):
     # Try to determine order by quick tags
     quick_tags, quick_tags_set = _get_quick_tags(quick_tags_text)
@@ -100,7 +105,7 @@ def update_image_tags(
     ]
     caption = ", ".join(output_tags)
 
-    if auto_save:
+    if auto_save_is_checked:
         save_caption(caption, caption_ext, image_file, images_dir)
 
     return caption
@@ -109,400 +114,207 @@ def update_image_tags(
 def import_tags_from_captions(
     images_dir, caption_ext, quick_tags_text, ignore_load_tags_word_count
 ):
-    """
-    Scans images directory for all available captions and loads all tags
-    under a specified word count into the quick tags box
-    """
-
-    def empty_return():
-        return gr.Text()
-
-    # Check for images_dir
-    if not images_dir:
-        msgbox("Image folder is missing...")
-        return empty_return()
-
-    if not os.path.exists(images_dir):
-        msgbox("Image folder does not exist...")
-        return empty_return()
+    if not images_dir or not os.path.exists(images_dir):
+        gr.Warning("Image folder is not set or does not exist. Please load images first.")
+        return gr.update()
 
     if not caption_ext:
-        msgbox("Please provide an extension for the caption files.")
-        return empty_return()
+        gr.Warning("Please provide an extension for the caption files.")
+        return gr.update()
 
     if quick_tags_text:
-        if not boolbox(
-            f"Are you sure you wish to overwrite the current quick tags?",
-            choices=("Yes", "No"),
-        ):
-            return empty_return()
+        if not boolbox("Are you sure you wish to overwrite the current quick tags?", choices=("Yes", "No")):
+            return gr.update()
 
-    images_list = os.listdir(images_dir)
-    image_files = [f for f in images_list if f.lower().endswith(IMAGE_EXTENSIONS)]
-
-    # Use a set for lookup but store order with list
     tags = []
     tags_set = set()
-    for image_file in image_files:
-        caption_file_path = _get_caption_path(image_file, images_dir, caption_ext)
-        if os.path.exists(caption_file_path):
-            with open(caption_file_path, "r", encoding="utf-8") as f:
-                caption = f.read()
-                for tag in caption.split(","):
-                    tag = tag.strip()
-                    tag_key = tag.lower()
-                    if not tag_key in tags_set:
-                        # Ignore extra spaces
-                        total_words = len(re.findall(r"\s+", tag)) + 1
-                        if total_words <= ignore_load_tags_word_count:
-                            tags.append(tag)
-                            tags_set.add(tag_key)
-
-    return ", ".join(tags)
-
-
-def load_images(images_dir, caption_ext, loaded_images_dir, page, max_page):
-    """
-    Triggered to load a new set of images from the folder to caption
-    This loads in the total expected image counts to be used by pagination
-    before running update_images
-    """
+    for entry in os.scandir(images_dir):
+        if entry.is_file() and entry.name.lower().endswith(IMAGE_EXTENSIONS):
+            caption_file_path = _get_caption_path(entry.name, images_dir, caption_ext)
+            if os.path.exists(caption_file_path):
+                with open(caption_file_path, "r", encoding="utf-8") as f:
+                    caption = f.read()
+                    for tag in caption.split(","):
+                        tag = tag.strip()
+                        if not tag:
+                            continue
+                        tag_key = tag.lower()
+                        if tag_key not in tags_set:
+                            if tag.count(" ") + 1 <= ignore_load_tags_word_count:
+                                tags.append(tag)
+                                tags_set.add(tag_key)
+
+    gr.Info(f"Imported {len(tags)} unique tags.")
+    return ", ".join(sorted(tags, key=str.lower))
+
+
+def load_images(images_dir, caption_ext):
+    def error_message(msg):
+        gr.Warning(msg)
+        return [None, None, 1, 1, gr.Markdown(f"⚠️ {msg}", visible=True)]
+
+    if not images_dir or not os.path.exists(images_dir):
+        return error_message("Image folder is missing or does not exist.")
+    if not caption_ext:
+        return error_message("Please provide an extension for the caption files.")
 
-    def empty_return():
-        return [loaded_images_dir, page, max_page]
+    image_files = sorted([f for f in os.listdir(images_dir) if f.lower().endswith(IMAGE_EXTENSIONS)])
 
-    # Check for images_dir
-    if not images_dir:
-        msgbox("Image folder is missing...")
-        return empty_return()
+    if not image_files:
+        return error_message("No images found in the folder.")
 
-    if not os.path.exists(images_dir):
-        msgbox("Image folder does not exist...")
-        return empty_return()
+    total_images = len(image_files)
+    max_pages = ceil(total_images / IMAGES_TO_SHOW)
 
-    if not caption_ext:
-        msgbox("Please provide an extension for the caption files.")
-        return empty_return()
+    info = f"✅ Loaded {total_images} images. {max_pages} pages total."
+    gr.Info(info)
 
-    # Load Images
-    images_list = os.listdir(images_dir)
-    total_images = len(
-        [True for f in images_list if f.lower().endswith(IMAGE_EXTENSIONS)]
-    )
-    return [images_dir, 1, ceil(total_images / IMAGES_TO_SHOW)]
+    return [image_files, images_dir, 1, max_pages, gr.Markdown(info, visible=True)]
 
 
 def update_images(
+    image_files,
     images_dir,
     caption_ext,
     quick_tags_text,
     page,
 ):
-    """
-    Updates the displayed images and captions from the current page and
-    image directory
-    """
-
-    # Load Images
-    images_list = os.listdir(images_dir)
-    image_files = [f for f in images_list if f.lower().endswith(IMAGE_EXTENSIONS)]
+    if not image_files or not images_dir:
+        empty_row = gr.Row(visible=False)
+        return [empty_row] * (IMAGES_TO_SHOW * 4 + 2)
 
-    # Quick tags
     quick_tags, quick_tags_set = _get_quick_tags(quick_tags_text or "")
 
-    # Display Images
-    rows = []
-    image_paths = []
-    captions = []
-    tag_checkbox_groups = []
-
+    outputs = []
     start_index = (int(page) - 1) * IMAGES_TO_SHOW
+
+    rows_update, files_update, paths_update, captions, tags_checks = [], [], [], [], []
+
     for i in range(IMAGES_TO_SHOW):
         image_index = start_index + i
-        show_row = image_index < len(image_files)
+        is_visible = image_index < len(image_files)
+
+        rows_update.append(gr.Row(visible=is_visible))
 
-        image_path = None
-        caption = ""
-        tag_checkboxes = None
-        if show_row:
+        image_file, image_path, caption = None, None, ""
+        if is_visible:
             image_file = image_files[image_index]
             image_path = os.path.join(images_dir, image_file)
-
             caption_file_path = _get_caption_path(image_file, images_dir, caption_ext)
-            if os.path.exists(caption_file_path):
+            if caption_file_path and os.path.exists(caption_file_path):
                 with open(caption_file_path, "r", encoding="utf-8") as f:
                     caption = f.read()
 
-        tag_checkboxes = _get_tag_checkbox_updates(caption, quick_tags, quick_tags_set)
-        rows.append(gr.Row(visible=show_row))
-        image_paths.append(image_path)
+        files_update.append(image_file)
+        paths_update.append(image_path)
         captions.append(caption)
-        tag_checkbox_groups.append(tag_checkboxes)
+        tags_checks.append(_get_tag_checkbox_updates(caption, quick_tags, quick_tags_set))
 
-    return (
-        rows
-        + image_paths
-        + image_paths
-        + captions
-        + tag_checkbox_groups
-        + [gr.Row(visible=True), gr.Row(visible=True)]
-    )
+    outputs.extend(rows_update)
+    outputs.extend(files_update)
+    outputs.extend(paths_update)
+    outputs.extend(captions)
+    outputs.extend(tags_checks)
+    outputs.extend([gr.Row(visible=True), gr.Row(visible=True)])
+
+    return outputs
 
 
 # Gradio UI
 def gradio_manual_caption_gui_tab(headless=False, default_images_dir=None):
     from .common_gui import create_refresh_button
 
-    default_images_dir = (
-        default_images_dir
-        if default_images_dir is not None
-        else os.path.join(scriptdir, "data")
-    )
-    current_images_dir = default_images_dir
+    default_images_dir = default_images_dir or os.path.join(scriptdir, "data")
+
+    def update_dir_list(path):
+        return gr.Dropdown(choices=[""] + list(list_dirs(path)))
 
-    # Function to list directories
-    def list_images_dirs(path):
-        # Allows list_images_dirs to modify current_images_dir outside of this function
-        nonlocal current_images_dir
-        current_images_dir = path
-        return list(list_dirs(path))
+    def render_pagination_with_logic(page, max_page):
+        with gr.Row(visible=False) as pagination_row:
+            gr.Button("◀ Prev").click(paginate, inputs=[page, max_page, gr.Number(-1, visible=False)], outputs=[page])
+            page_count = gr.Text("Page 1 / 1", show_label=False, interactive=False, text_align="center")
+            page_goto_text = gr.Textbox(show_label=False, placeholder="Go to page...", container=False, scale=1)
+            gr.Button("Next ▶").click(paginate, inputs=[page, max_page, gr.Number(1, visible=False)], outputs=[page])
+            page_goto_text.submit(paginate_go, inputs=[page_goto_text, max_page], outputs=[page])
+        return pagination_row, page_count
 
     with gr.Tab("Manual Captioning"):
         gr.Markdown("This utility allows quick captioning and tagging of images.")
-        page = gr.Number(value=-1, visible=False)
+
+        image_files_state = gr.State([])
+
+        info_box = gr.Markdown(visible=False)
+        page = gr.State(value=1)
         max_page = gr.Number(value=1, visible=False)
         loaded_images_dir = gr.Text(visible=False)
-        with gr.Group(), gr.Row():
-            images_dir = gr.Dropdown(
-                label="Image folder to caption (containing the images to caption)",
-                choices=[""] + list_images_dirs(default_images_dir),
-                value="",
-                interactive=True,
-                allow_custom_value=True,
-            )
-            create_refresh_button(
-                images_dir,
-                lambda: None,
-                lambda: {"choices": list_images_dirs(current_images_dir)},
-                "open_folder_small",
-            )
-            folder_button = gr.Button(
-                "📂",
-                elem_id="open_folder_small",
-                elem_classes=["tool"],
-                visible=(not headless),
-            )
-            folder_button.click(
-                get_folder_path,
-                outputs=images_dir,
-                show_progress=False,
-            )
-            load_images_button = gr.Button("Load", elem_id="open_folder")
-            caption_ext = gr.Dropdown(
-                label="Caption file extension",
-                choices=[".cap", ".caption", ".txt"],
-                value=".txt",
-                interactive=True,
-                allow_custom_value=True,
-            )
-            auto_save = gr.Checkbox(
-                label="Autosave", info="Options", value=True, interactive=True
-            )
-
-            images_dir.change(
-                fn=lambda path: gr.Dropdown(choices=[""] + list_images_dirs(path)),
-                inputs=images_dir,
-                outputs=images_dir,
-                show_progress=False,
-            )
-
-        # Caption Section
-        with gr.Group(), gr.Row():
-            quick_tags_text = gr.Textbox(
-                label="Quick Tags",
-                placeholder="Comma separated list of tags",
-                interactive=True,
-            )
-            import_tags_button = gr.Button("Import", elem_id="open_folder")
-            ignore_load_tags_word_count = gr.Slider(
-                minimum=1,
-                maximum=100,
-                value=3,
-                step=1,
-                label="Ignore Imported Tags Above Word Count",
-                interactive=True,
-            )
-
-        # Next/Prev section generator
-        def render_pagination():
-            gr.Button("< Prev", elem_id="open_folder").click(
-                paginate,
-                inputs=[page, max_page, gr.Number(value=-1, visible=False)],
-                outputs=[page],
-            )
-            page_count = gr.Label("Page 1", label="Page")
-            page_goto_text = gr.Textbox(
-                label="Goto page",
-                placeholder="Page Number",
-                interactive=True,
-            )
-            gr.Button("Go >", elem_id="open_folder").click(
-                paginate_go,
-                inputs=[page_goto_text, max_page],
-                outputs=[page],
-            )
-            gr.Button("Next >", elem_id="open_folder").click(
-                paginate,
-                inputs=[page, max_page, gr.Number(value=1, visible=False)],
-                outputs=[page],
-            )
-            return page_count
-
-        with gr.Row(visible=False) as pagination_row1:
-            page_count1 = render_pagination()
-
-        # Images section
-        image_rows = []
-        image_files = []
-        image_images = []
-        image_caption_texts = []
-        image_tag_checks = []
-        save_buttons = []
-        for _ in range(IMAGES_TO_SHOW):
+
+        with gr.Group():
+            with gr.Row():
+                images_dir = gr.Dropdown(label="Image folder to caption", choices=[""] + list(list_dirs(default_images_dir)), value="", interactive=True, allow_custom_value=True)
+                create_refresh_button(images_dir, lambda: None, lambda: {"choices": list(list_dirs(images_dir.value or default_images_dir))}, "open_folder_small")
+                gr.Button("📂", elem_id="open_folder_small", elem_classes=["tool"], visible=not headless).click(get_folder_path, outputs=images_dir, show_progress=False)
+
+
+            with gr.Row():
+                caption_ext = gr.Dropdown(label="Caption file extension", choices=[".cap", ".caption", ".txt"], value=".txt", interactive=True, allow_custom_value=True)
+                auto_save = gr.Checkbox(label="Autosave", value=True, interactive=True)
+                
+            with gr.Row():
+                load_images_button = gr.Button("Load Images", variant="primary")
+
+            images_dir.change(update_dir_list, inputs=images_dir, outputs=images_dir, show_progress=False)
+
+        with gr.Group():
+            quick_tags_text = gr.Textbox(label="Quick Tags", placeholder="Comma separated list of tags", interactive=True)
+            with gr.Row():
+                ignore_load_tags_word_count = gr.Slider(minimum=1, maximum=100, value=10, step=1, label="Ignore Imported Tags Above Word Count", interactive=True, scale=2)
+            
+            with gr.Row():
+                import_tags_button = gr.Button("Import tags from captions", scale=1)
+
+        pagination_row1, page_count1 = render_pagination_with_logic(page, max_page)
+
+        image_rows, image_files, image_images, image_caption_texts, image_tag_checks, save_buttons = [], [], [], [], [], []
+        for i in range(IMAGES_TO_SHOW):
             with gr.Row(visible=False) as row:
                 image_file = gr.Text(visible=False)
+                with gr.Column():
+                    image_image = gr.Image(type="filepath", label=f"Image {i+1}")
+                with gr.Column(scale=2):
+                    image_caption_text = gr.TextArea(label="Captions", placeholder="Input captions for image", interactive=True)
+                    tag_checkboxes = gr.CheckboxGroup([], label="Tags", interactive=True)
+                with gr.Column(min_width=40):
+                    save_button = gr.Button("💾", elem_id="save_button", visible=False)
+
+                image_rows.append(row)
                 image_files.append(image_file)
-                image_image = gr.Image(type="filepath")
                 image_images.append(image_image)
-                image_caption_text = gr.TextArea(
-                    label="Captions",
-                    placeholder="Input captions",
-                    interactive=True,
-                )
                 image_caption_texts.append(image_caption_text)
-                tag_checkboxes = gr.CheckboxGroup([], label="Tags", interactive=True)
-                save_button = gr.Button(
-                    "💾",
-                    elem_id="open_folder_small",
-                    elem_classes=["tool"],
-                    visible=False,
-                )
+                image_tag_checks.append(tag_checkboxes)
                 save_buttons.append(save_button)
 
-                # Caption text change
-                image_caption_text.input(
-                    update_image_caption,
-                    inputs=[
-                        quick_tags_text,
-                        image_caption_text,
-                        image_file,
-                        loaded_images_dir,
-                        caption_ext,
-                        auto_save,
-                    ],
-                    outputs=tag_checkboxes,
-                )
-
-                # Quick tag check
-                tag_checkboxes.input(
-                    update_image_tags,
-                    inputs=[
-                        quick_tags_text,
-                        tag_checkboxes,
-                        image_file,
-                        loaded_images_dir,
-                        caption_ext,
-                        auto_save,
-                    ],
-                    outputs=[image_caption_text],
-                )
-
-                # Save Button
-                save_button.click(
-                    save_caption,
-                    inputs=[
-                        image_caption_text,
-                        caption_ext,
-                        image_file,
-                        images_dir,
-                    ],
-                )
+                image_caption_text.input(update_image_caption, inputs=[quick_tags_text, image_caption_text, image_file, loaded_images_dir, caption_ext, auto_save], outputs=tag_checkboxes)
+                tag_checkboxes.input(update_image_tags, inputs=[quick_tags_text, tag_checkboxes, image_file, loaded_images_dir, caption_ext, auto_save], outputs=[image_caption_text])
+                save_button.click(save_caption, inputs=[image_caption_text, caption_ext, image_file, loaded_images_dir], outputs=info_box)
 
-                image_tag_checks.append(tag_checkboxes)
-                image_rows.append(row)
+        pagination_row2, page_count2 = render_pagination_with_logic(page, max_page)
 
-        # Next/Prev Section
-        with gr.Row(visible=False) as pagination_row2:
-            page_count2 = render_pagination()
+        quick_tags_text.change(update_quick_tags, inputs=[quick_tags_text] + image_caption_texts, outputs=image_tag_checks)
+        import_tags_button.click(import_tags_from_captions, inputs=[loaded_images_dir, caption_ext, quick_tags_text, ignore_load_tags_word_count], outputs=quick_tags_text)
 
-        # Quick tag text update
-        quick_tags_text.change(
-            update_quick_tags,
-            inputs=[quick_tags_text] + image_caption_texts,
-            outputs=image_tag_checks,
-        )
+        load_images_outputs = [image_files_state, loaded_images_dir, page, max_page, info_box]
+        load_images_button.click(load_images, inputs=[images_dir, caption_ext], outputs=load_images_outputs)
 
-        # Import tags button
-        import_tags_button.click(
-            import_tags_from_captions,
-            inputs=[
-                loaded_images_dir,
-                caption_ext,
-                quick_tags_text,
-                ignore_load_tags_word_count,
-            ],
-            outputs=quick_tags_text,
+        update_trigger_inputs = [image_files_state, loaded_images_dir, caption_ext, quick_tags_text, page]
+        update_outputs = (
+            image_rows + image_files + image_images +
+            image_caption_texts + image_tag_checks + [pagination_row1, pagination_row2]
         )
 
-        # Load Images button
-        load_images_button.click(
-            load_images,
-            inputs=[
-                images_dir,
-                caption_ext,
-                loaded_images_dir,
-                page,
-                max_page,
-            ],
-            outputs=[loaded_images_dir, page, max_page],
-        )
+        page.change(update_images, inputs=update_trigger_inputs, outputs=update_outputs, show_progress=False)
+        image_files_state.change(update_images, inputs=update_trigger_inputs, outputs=update_outputs, show_progress=False)
 
-        # Update images shown when the update key changes
-        # This allows us to trigger a change from multiple
-        # sources (page, image_dir)
-        image_update_key = gr.Text(visible=False)
-        image_update_key.change(
-            update_images,
-            inputs=[loaded_images_dir, caption_ext, quick_tags_text, page],
-            outputs=image_rows
-            + image_files
-            + image_images
-            + image_caption_texts
-            + image_tag_checks
-            + [pagination_row1, pagination_row2],
-            show_progress=False,
-        )
-        # Update the key on page and image dir change
-        listener_kwargs = {
-            "fn": lambda p, i: f"{p}-{i}",
-            "inputs": [page, loaded_images_dir],
-            "outputs": image_update_key,
-        }
-        page.change(**listener_kwargs)
-        loaded_images_dir.change(**listener_kwargs)
-
-        # Save buttons visibility
-        # (on auto-save on/off)
-        auto_save.change(
-            lambda auto_save: [gr.Button(visible=not auto_save)] * IMAGES_TO_SHOW,
-            inputs=auto_save,
-            outputs=save_buttons,
-        )
-
-        # Page Count
-        page.change(
-            lambda page, max_page: [f"Page {int(page)} / {int(max_page)}"] * 2,
-            inputs=[page, max_page],
-            outputs=[page_count1, page_count2],
-            show_progress=False,
-        )
+        auto_save.change(lambda is_auto_save: [gr.Button(visible=not is_auto_save)] * IMAGES_TO_SHOW, inputs=auto_save, outputs=save_buttons)
+        page.change(lambda p, m: (f"Page {int(p)} / {int(m)}", f"Page {int(p)} / {int(m)}"), inputs=[page, max_page], outputs=[page_count1, page_count2], show_progress=False)
+        image_files_state.change(lambda p, m: (f"Page {int(p)} / {int(m)}", f"Page {int(p)} / {int(m)}"), inputs=[page, max_page], outputs=[page_count1, page_count2], show_progress=False)
+        max_page.change(lambda p, m: (f"Page {int(p)} / {int(m)}", f"Page {int(p)} / {int(m)}"), inputs=[page, max_page], outputs=[page_count1, page_count2], show_progress=False)
diff --git a/kohya_gui/utilities.py b/kohya_gui/utilities.py
index 143b033d5..ec76e572c 100644
--- a/kohya_gui/utilities.py
+++ b/kohya_gui/utilities.py
@@ -26,6 +26,11 @@ def utilities_tab(
         gradio_git_caption_gui_tab(headless=headless)
         gradio_wd14_caption_gui_tab(headless=headless, config=config)
         gradio_manual_caption_gui_tab(headless=headless)
+        from kohya_gui.kontext_manual_caption_gui import (
+            gradio_kontext_manual_caption_gui_tab,
+        )
+
+        gradio_kontext_manual_caption_gui_tab(headless=headless)
     gradio_convert_model_tab(headless=headless)
     gradio_group_images_gui_tab(headless=headless)
 
diff --git a/kohya_gui/wd14_caption_gui.py b/kohya_gui/wd14_caption_gui.py
index 3456f70be..f8a6c7513 100644
--- a/kohya_gui/wd14_caption_gui.py
+++ b/kohya_gui/wd14_caption_gui.py
@@ -40,6 +40,7 @@ def caption_images(
     use_rating_tags_as_last_tag: bool,
     remove_underscore: bool,
     thresh: float,
+    sort_tags_by_confidence: bool,
 ) -> None:
     # Check for images_dir_input
     if train_data_dir == "":
@@ -113,6 +114,8 @@ def caption_images(
         run_cmd.append("--use_rating_tags")
     if use_rating_tags_as_last_tag:
         run_cmd.append("--use_rating_tags_as_last_tag")
+    if sort_tags_by_confidence:
+        run_cmd.append("--sort_tags_by_confidence")
 
     # Add the directory containing the training data
     run_cmd.append(rf"{train_data_dir}")
@@ -317,6 +320,11 @@ def list_train_dirs(path):
                 value=config.get("wd14_caption.frequency_tags", True),
                 info="Show frequency of tags for images.",
             )
+            sort_tags_by_confidence = gr.Checkbox(
+                label="Sort tags by confidence",
+                value=config.get("wd14_caption.sort_tags_by_confidence", False),
+                info="Sort output tags by confidence score (highest first), matching the WaifuDiffusion online tagger order.",
+            )
 
         with gr.Row():
             thresh = gr.Slider(
@@ -396,6 +404,7 @@ def repo_id_changes(repo_id, onnx):
                 use_rating_tags_as_last_tag,
                 remove_underscore,
                 thresh,
+                sort_tags_by_confidence,
             ],
             show_progress=False,
         )
diff --git a/kohya_ss_colab.ipynb b/kohya_ss_colab.ipynb
new file mode 100644
index 000000000..0002aee74
--- /dev/null
+++ b/kohya_ss_colab.ipynb
@@ -0,0 +1,103 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github"
+      },
+      "source": [
+        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/bmaltais/kohya_ss/blob/feature/update-gui-uv-for-colab-parity/kohya_ss_colab.ipynb)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "#@title Train with Kohya's Stable Diffusion Trainers\n",
+        "%cd /content\n",
+        "\n",
+        "import os\n",
+        "from google.colab import drive\n",
+        "drive.mount('/content/drive')\n",
+        "\n",
+        "%cd /content\n",
+        "!git clone -b dev https://github.com/bmaltais/kohya_ss\n",
+        "%cd /content/kohya_ss\n",
+        "!git checkout dev\n",
+        "\n",
+        "os.environ[\"TERM\"] = \"dumb\"\n",
+        "!./gui-uv.sh --quiet --share --headless"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "#@title Convert Safetensors to Diffusers\n",
+        "from_safetensors_url = '' #@param {type:\"string\"}\n",
+        "!wget -q https://raw.githubusercontent.com/huggingface/diffusers/v0.17.1/scripts/convert_original_stable_diffusion_to_diffusers.py\n",
+        "!wget {from_safetensors_url} -O /content/model.safetensors\n",
+        "!python3 convert_original_stable_diffusion_to_diffusers.py --half --from_safetensors --checkpoint_path model.safetensors --dump_path /content/model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "#@title Push to HF.co\n",
+        "\n",
+        "import os\n",
+        "from huggingface_hub import create_repo, upload_folder\n",
+        "\n",
+        "hf_token = 'HF_WRITE_TOKEN' #@param {type:\"string\"}\n",
+        "repo_id = 'username/reponame' #@param {type:\"string\"}\n",
+        "commit_message = '\\u2764' #@param {type:\"string\"}\n",
+        "create_repo(repo_id, private=True, token=hf_token)\n",
+        "model_path = '/content/train/model' #@param {type:\"string\"}\n",
+        "upload_folder(folder_path=f'{model_path}', path_in_repo='', repo_id=repo_id, commit_message=commit_message, token=hf_token)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "#@title Push to DagsHub.com\n",
+        "\n",
+        "!pip -q install dagshub\n",
+        "from dagshub.upload import Repo, create_repo\n",
+        "\n",
+        "repo_id = 'reponame' #@param {type:\"string\"}\n",
+        "org_name = 'orgname' #@param {type:\"string\"}\n",
+        "commit_message = '\\u2764' #@param {type:\"string\"}\n",
+        "create_repo(f\"{repo_id}\", org_name=f\"{org_name}\")\n",
+        "repo = Repo(f\"{org_name}\", f\"{repo_id}\")\n",
+        "model_path = '/content/train/model' #@param {type:\"string\"}\n",
+        "repo.upload(\"/content/models\", remote_path=\"data\", commit_message=f\"{commit_message}\", force=True)"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "gpuType": "T4",
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
diff --git a/local_train.py b/local_train.py
new file mode 100644
index 000000000..b4563549b
--- /dev/null
+++ b/local_train.py
@@ -0,0 +1,81 @@
+import os
+import subprocess
+
+def run_command(command):
+    """Runs a shell command and prints its output."""
+    print(f"Running command: {command}")
+    process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
+    for line in process.stdout:
+        print(line, end="")
+    process.wait()
+    if process.returncode != 0:
+        raise Exception(f"Command failed with return code {process.returncode}: {command}")
+
+def main():
+    print("Starting local training setup...")
+
+    # Target directory for cloning and operations, can be made configurable
+    base_dir = "kohya_local_setup"
+    os.makedirs(base_dir, exist_ok=True)
+    original_dir = os.getcwd()
+    os.chdir(base_dir)
+
+    print(f"Working directory: {os.getcwd()}")
+
+    # 1. Install dependencies (This should ideally be handled by gui-uv.sh using requirements.txt)
+    # For now, we'll list them. These were in the notebook:
+    # dadaptation==3.1 diffusers[torch]==0.17.1 easygui==0.98.3 einops==0.6.0
+    # fairscale==0.4.13 ftfy==6.1.1 gradio==3.36.1 huggingface-hub==0.14.1
+    # lion-pytorch==0.0.6 lycoris_lora==1.8.0.dev6 open-clip-torch==2.20.0
+    # prodigyopt==1.0 pytorch-lightning==1.9.0 safetensors==0.3.1 timm==0.6.12
+    # tk==0.1.0 transformers==4.30.2 voluptuous==0.13.1 wandb==0.15.0
+    # xformers==0.0.20 omegaconf
+    print("Ensure all dependencies are installed via requirements.txt by gui-uv.sh")
+
+    # 2. Clone bitsandbytes
+    if not os.path.exists("bitsandbytes"):
+        run_command("git clone -b 0.41.0 https://github.com/TimDettmers/bitsandbytes")
+    else:
+        print("bitsandbytes directory already exists. Skipping clone.")
+
+    # 3. Build and install bitsandbytes
+    # Note: This might need specific CUDA versions. The notebook used CUDA_VERSION=118
+    # This part is tricky for a generic local setup and might need user intervention or
+    # pre-compiled versions if not handled by the main gui-uv.sh environment.
+    # For now, we are replicating the notebook's steps.
+    # It's better if bitsandbytes is installed as a pip package if possible.
+    print("Building and installing bitsandbytes...")
+    os.chdir("bitsandbytes")
+    run_command("make cuda11x") # This assumes CUDA 11.8 development toolkit is available
+    run_command("python setup.py install")
+    os.chdir("..") # Back to base_dir
+
+    # 4. Clone kohya_ss
+    if not os.path.exists("kohya_ss"):
+        run_command("git clone -b v1.0 https://github.com/camenduru/kohya_ss")
+    else:
+        print("kohya_ss directory already exists. Skipping clone.")
+        # Optionally, add a git pull here to update
+        # os.chdir("kohya_ss")
+        # run_command("git pull")
+        # os.chdir("..")
+
+
+    # 5. Launch Kohya GUI
+    print("Launching Kohya GUI...")
+    os.chdir("kohya_ss")
+    # The notebook uses --share --headless.
+    # --share might not be needed or desired for local use.
+    # --headless might be for Colab's environment. For local GUI, we might not want headless.
+    # We need to check how gui-uv.sh expects to launch this.
+    # For now, using the notebook's command but this will likely need adjustment.
+    try:
+        run_command("python kohya_gui.py --headless")
+    except Exception as e:
+        print(f"Error launching Kohya GUI: {e}")
+        print("Please ensure all dependencies, including CUDA and xformers, are correctly installed.")
+    finally:
+        os.chdir(original_dir) # Change back to the original directory
+
+if __name__ == "__main__":
+    main()
diff --git a/pyproject.toml b/pyproject.toml
index f54989684..e95b95efb 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -31,6 +31,7 @@ dependencies = [
     "prodigy-plus-schedule-free==1.8.0",
     "prodigyopt==1.1.2",
     "protobuf==3.20.3",
+    "py-cpuinfo",
     "pytorch-lightning==1.9.0",
     "pytorch-optimizer==3.5.0",
     "rich>=13.7.1",
diff --git a/requirements.txt b/requirements.txt
index 0e2bd2882..85442c6eb 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
+py-cpuinfo
 accelerate>=1.7.0
 aiofiles==23.2.1
 altair==4.2.2
diff --git a/requirements_linux_ipex.txt b/requirements_linux_ipex.txt
index 8a9187cdc..177ec7737 100644
--- a/requirements_linux_ipex.txt
+++ b/requirements_linux_ipex.txt
@@ -1,19 +1,22 @@
 # Custom index URL for specific packages
+--extra-index-url https://download.pytorch.org/whl/xpu
 --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/us/
 
-torch==2.3.1+cxx11.abi
-torchvision==0.18.1+cxx11.abi
-intel-extension-for-pytorch==2.3.110+xpu
-oneccl_bind_pt==2.3.100+xpu
+torch==2.7.0+xpu
+torchvision==0.22.0+xpu
+intel-extension-for-pytorch==2.7.10+xpu
+oneccl_bind_pt==2.7.0+xpu
 
 tensorboard==2.15.2
-tensorflow==2.15.1
-intel-extension-for-tensorflow[xpu]==2.15.0.1
+tensorflow==2.15.0
+intel-extension-for-tensorflow[xpu]==2.15.0.2
+onnxruntime==1.22.1
 onnxruntime-openvino==1.22.0
 
-mkl==2024.2.1
-mkl-dpcpp==2024.2.1
-oneccl-devel==2021.13.1
-impi-devel==2021.13.1
+mkl==2025.0.1
+mkl-dpcpp==2025.0.1
+oneccl-devel==2021.14.1
+impi-devel==2021.14.1
 
 -r requirements.txt
+
diff --git a/requirements_macos_arm64.txt b/requirements_macos_arm64.txt
index c495f9d02..26f590852 100644
--- a/requirements_macos_arm64.txt
+++ b/requirements_macos_arm64.txt
@@ -1,7 +1,6 @@
---extra-index-url https://download.pytorch.org/whl/nightly/cpu
+--pre --extra-index-url https://download.pytorch.org/whl/nightly/cpu
 torch==2.8.0.*
 torchvision==0.22.*
-xformers==0.0.29.* 
 git+https://github.com/bitsandbytes-foundation/bitsandbytes.git/#0.45.5
 tensorflow-macos 
 tensorflow-metal
diff --git a/sd-scripts b/sd-scripts
index 3e6935a07..2e0fcc50c 160000
--- a/sd-scripts
+++ b/sd-scripts
@@ -1 +1 @@
-Subproject commit 3e6935a07edcb944407840ef74fcaf6fcad352f7
+Subproject commit 2e0fcc50cb347323b4374bfd05784587c5fa7f2d
diff --git a/test/test_allowed_paths.py b/test/test_allowed_paths.py
new file mode 100644
index 000000000..a99b8a7b3
--- /dev/null
+++ b/test/test_allowed_paths.py
@@ -0,0 +1,59 @@
+import os
+import tempfile
+import unittest
+import subprocess
+from pathlib import Path
+
+class TestAllowedPaths(unittest.TestCase):
+    def setUp(self):
+        print("Setting up test...")
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.temp_dir_path = Path(self.temp_dir.name)
+        self.dummy_file = self.temp_dir_path / "dummy.txt"
+        with open(self.dummy_file, "w") as f:
+            f.write("dummy content")
+
+        self.config_file = self.temp_dir_path / "config.toml"
+        with open(self.config_file, "w") as f:
+            f.write(f'''
+[server]
+allowed_paths = ["{self.temp_dir.name}"]
+''')
+        print("Setup complete.")
+
+    def tearDown(self):
+        print("Tearing down test...")
+        self.temp_dir.cleanup()
+        print("Teardown complete.")
+
+    def test_allowed_paths(self):
+        print("Running test_allowed_paths...")
+        # Run the gui with the new config and check if it can access the dummy file
+        process = subprocess.Popen(
+            [
+                "python",
+                "kohya_gui.py",
+                "--config",
+                str(self.config_file),
+                "--headless",
+            ],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+
+        print("Process started.")
+        # Give the server some time to start
+        try:
+            stdout, stderr = process.communicate(timeout=10)
+        except subprocess.TimeoutExpired:
+            process.kill()
+            stdout, stderr = process.communicate()
+
+        print(f"Stdout: {stdout.decode()}")
+        print(f"Stderr: {stderr.decode()}")
+        # Check if there are any errors in the stderr
+        self.assertNotIn("InvalidPathError", stderr.decode())
+        print("Test complete.")
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_tensorboard_visibility.py b/tests/test_tensorboard_visibility.py
new file mode 100644
index 000000000..83ccc53b6
--- /dev/null
+++ b/tests/test_tensorboard_visibility.py
@@ -0,0 +1,34 @@
+import unittest
+from unittest.mock import patch
+import importlib
+
+# Since we are modifying an existing file, we need to reload it
+import kohya_gui.class_tensorboard
+importlib.reload(kohya_gui.class_tensorboard)
+
+class TestTensorboardVisibility(unittest.TestCase):
+
+    @patch('shutil.which', return_value='/usr/bin/tensorboard')
+    @patch('cpuinfo.get_cpu_info', return_value={'flags': ['avx']})
+    def test_tensorboard_visibility_when_tensorboard_and_avx_are_present(self, mock_cpuinfo, mock_which):
+        importlib.reload(kohya_gui.class_tensorboard)
+        self.assertTrue(kohya_gui.class_tensorboard.visibility)
+
+    @patch('shutil.which', return_value=None)
+    @patch('cpuinfo.get_cpu_info', return_value={'flags': ['avx']})
+    def test_tensorboard_visibility_when_tensorboard_is_absent(self, mock_cpuinfo, mock_which):
+        importlib.reload(kohya_gui.class_tensorboard)
+        self.assertFalse(kohya_gui.class_tensorboard.visibility)
+
+    @patch('shutil.which', return_value='/usr/bin/tensorboard')
+    @patch('cpuinfo.get_cpu_info', return_value={'flags': ['sse']})
+    def test_tensorboard_visibility_when_avx_is_absent(self, mock_cpuinfo, mock_which):
+        importlib.reload(kohya_gui.class_tensorboard)
+        self.assertFalse(kohya_gui.class_tensorboard.visibility)
+
+    @patch('cpuinfo.get_cpu_info', side_effect=Exception)
+    def test_check_avx_support_exception(self, mock_cpuinfo):
+        self.assertFalse(kohya_gui.class_tensorboard.check_avx_support())
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/uv.lock b/uv.lock
index a80eea3e3..1af1a8022 100644
--- a/uv.lock
+++ b/uv.lock
@@ -859,6 +859,7 @@ dependencies = [
     { name = "prodigy-plus-schedule-free" },
     { name = "prodigyopt" },
     { name = "protobuf" },
+    { name = "py-cpuinfo" },
     { name = "pytorch-lightning" },
     { name = "pytorch-optimizer" },
     { name = "rich" },
@@ -914,6 +915,7 @@ requires-dist = [
     { name = "prodigy-plus-schedule-free", specifier = "==1.8.0" },
     { name = "prodigyopt", specifier = "==1.1.2" },
     { name = "protobuf", specifier = "==3.20.3" },
+    { name = "py-cpuinfo" },
     { name = "pytorch-lightning", specifier = "==1.9.0" },
     { name = "pytorch-optimizer", specifier = "==3.5.0" },
     { name = "rich", specifier = ">=13.7.1" },
@@ -1698,6 +1700,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/50/1b/6921afe68c74868b4c9fa424dad3be35b095e16687989ebbb50ce4fceb7c/psutil-7.0.0-cp37-abi3-win_amd64.whl", hash = "sha256:4cf3d4eb1aa9b348dec30105c55cd9b7d4629285735a102beb4441e38db90553", size = 244885, upload-time = "2025-02-13T21:54:37.486Z" },
 ]
 
+[[package]]
+name = "py-cpuinfo"
+version = "9.0.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/37/a8/d832f7293ebb21690860d2e01d8115e5ff6f2ae8bbdc953f0eb0fa4bd2c7/py-cpuinfo-9.0.0.tar.gz", hash = "sha256:3cdbbf3fac90dc6f118bfd64384f309edeadd902d7c8fb17f02ffa1fc3f49690", size = 104716, upload-time = "2022-10-25T20:38:06.303Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/e0/a9/023730ba63db1e494a271cb018dcd361bd2c917ba7004c3e49d5daf795a2/py_cpuinfo-9.0.0-py3-none-any.whl", hash = "sha256:859625bc251f64e21f077d099d4162689c762b5d6a4c3c97553d56241c9674d5", size = 22335, upload-time = "2022-10-25T20:38:27.636Z" },
+]
+
 [[package]]
 name = "pydantic"
 version = "2.11.7"