Sd webui tab extension support

install.py

@@ -0,0 +1,61 @@
+import os
+import platform
+import sys
+
+from packaging import version
+from pathlib import Path
+from typing import Tuple, Optional
+
+import launch
+from launch import is_installed, run, run_pip
+import importlib.metadata
+
+try:
+    skip_install = getattr(launch.args, "skip_install")
+except Exception:
+    skip_install = getattr(launch, "skip_install", False)
+
+python = sys.executable
+
+
+def comparable_version(version: str) -> Tuple:
+    return tuple(version.split("."))
+
+
+def get_installed_version(package: str) -> Optional[str]:
+    try:
+        return importlib.metadata.version(package)
+    except Exception:
+        return None
+
+
+def install():
+    req_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), "requirements-sdwebui.txt")
+    if os.path.exists(req_file):
+        mainpackage = 'sd-scripts'
+        with open(req_file) as file:
+            for package in file:
+                try:
+                    package = package.strip()
+                    if package[0] in ["#", "-"]:
+                        continue
+
+                    if '==' in package:
+                        package_name, package_version = package.split('==')
+                        installed_version = get_installed_version(package_name)
+                        if installed_version != package_version:
+                            run_pip(f"install -U {package}", f"{mainpackage} requirement: changing {package_name} version from {installed_version} to {package_version}")
+                    elif '>=' in package:
+                        package_name, package_version = package.split('>=')
+                        installed_version = get_installed_version(package_name)
+                        if not installed_version or comparable_version(installed_version) < comparable_version(package_version):
+                            run_pip(f"install -U {package}", f"{mainpackage} requirement: changing {package_name} version from {installed_version} to {package_version}")
+                    elif not is_installed(package):
+                        run_pip(f"install {package}", f"{mainpackage} requirement: {package}")
+                except Exception as e:
+                    print(f"Error: {e}")
+                    print(f'Warning: Failed to install {package}, some preprocessors may not work.')
+
+
+if not skip_install:
+    install()

kohya_gui/__init__.py

@@ -0,0 +1 @@
+"""empty"""

kohya_gui/basic_caption_gui.py

@@ -0,0 +1,167 @@
+import gradio as gr
+from easygui import msgbox
+import subprocess
+from .common_gui import get_folder_path, add_pre_postfix, find_replace
+import os
+
+from .custom_logging import setup_logging
+
+# Set up logging
+log = setup_logging()
+
+try:
+    from modules import scripts, paths
+
+    topdir = paths.script_path
+    scriptdir = scripts.basedir()
+except Exception:
+    topdir = "."
+    scriptdir = "."
+
+
+def caption_images(
+    caption_text,
+    images_dir,
+    overwrite,
+    caption_ext,
+    prefix,
+    postfix,
+    find_text,
+    replace_text,
+):
+    # Check if images_dir is provided
+    if not images_dir:
+        msgbox(
+            'Image folder is missing. Please provide the directory containing the images to caption.'
+        )
+        return
+
+    # Check if caption_ext is provided
+    if not caption_ext:
+        msgbox('Please provide an extension for the caption files.')
+        return
+
+    if caption_text:
+        log.info(f'Captioning files in {images_dir} with {caption_text}...')
+
+        # Build the command to run caption.py
+        run_cmd = fr'python "{scriptdir}/tools/caption.py"'
+        run_cmd += f' --caption_text="{caption_text}"'
+
+        # Add optional flags to the command
+        if overwrite:
+            run_cmd += f' --overwrite'
+        if caption_ext:
+            run_cmd += f' --caption_file_ext="{caption_ext}"'
+
+        run_cmd += f' "{images_dir}"'
+
+        log.info(run_cmd)
+
+        # Run the command based on the operating system
+        if os.name == 'posix':
+            os.system(run_cmd)
+        else:
+            subprocess.run(run_cmd)
+
+    # Check if overwrite option is enabled
+    if overwrite:
+        if prefix or postfix:
+            # Add prefix and postfix to caption files
+            add_pre_postfix(
+                folder=images_dir,
+                caption_file_ext=caption_ext,
+                prefix=prefix,
+                postfix=postfix,
+            )
+        if find_text:
+            # Find and replace text in caption files
+            find_replace(
+                folder_path=images_dir,
+                caption_file_ext=caption_ext,
+                search_text=find_text,
+                replace_text=replace_text,
+            )
+    else:
+        if prefix or postfix:
+            # Show a message if modification is not possible without overwrite option enabled
+            msgbox(
+                'Could not modify caption files with requested change because the "Overwrite existing captions in folder" option is not selected.'
+            )
+
+    log.info('Captioning done.')
+
+
+# Gradio UI
+def gradio_basic_caption_gui_tab(headless=False):
+    with gr.Tab('Basic Captioning'):
+        gr.Markdown(
+            'This utility allows you to create simple caption files for each image in a folder.'
+        )
+        with gr.Row():
+            images_dir = gr.Textbox(
+                label='Image folder to caption',
+                placeholder='Directory containing the images to caption',
+                interactive=True,
+            )
+            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,
+            )
+            caption_ext = gr.Textbox(
+                label='Caption file extension',
+                placeholder='Extension for caption file (e.g., .caption, .txt)',
+                value='.txt',
+                interactive=True,
+            )
+            overwrite = gr.Checkbox(
+                label='Overwrite existing captions in folder',
+                interactive=True,
+                value=False,
+            )
+        with gr.Row():
+            prefix = gr.Textbox(
+                label='Prefix to add to caption',
+                placeholder='(Optional)',
+                interactive=True,
+            )
+            caption_text = gr.Textbox(
+                label='Caption text',
+                placeholder='e.g., "by some artist". Leave empty if you only want to add a prefix or postfix.',
+                interactive=True,
+            )
+            postfix = gr.Textbox(
+                label='Postfix to add to caption',
+                placeholder='(Optional)',
+                interactive=True,
+            )
+        with gr.Row():
+            find_text = gr.Textbox(
+                label='Find text',
+                placeholder='e.g., "by some artist". Leave empty if you only want to add a prefix or postfix.',
+                interactive=True,
+            )
+            replace_text = gr.Textbox(
+                label='Replacement text',
+                placeholder='e.g., "by some artist". Leave empty if you want to replace with nothing.',
+                interactive=True,
+            )
+            caption_button = gr.Button('Caption images')
+            caption_button.click(
+                caption_images,
+                inputs=[
+                    caption_text,
+                    images_dir,
+                    overwrite,
+                    caption_ext,
+                    prefix,
+                    postfix,
+                    find_text,
+                    replace_text,
+                ],
+                show_progress=False,
+            )

kohya_gui/blip_caption_gui.py

@@ -0,0 +1,160 @@
+import gradio as gr
+from easygui import msgbox
+import subprocess
+import os
+from .common_gui import get_folder_path, add_pre_postfix
+from .custom_logging import setup_logging
+
+# Set up logging
+log = setup_logging()
+
+try:
+    from modules import scripts, paths
+
+    topdir = paths.script_path
+    scriptdir = scripts.basedir()
+except Exception:
+    topdir = "."
+    scriptdir = "."
+
+PYTHON = 'python3' if os.name == 'posix' else fr'{topdir}/venv/Scripts/python.exe'
+
+
+def caption_images(
+    train_data_dir,
+    caption_file_ext,
+    batch_size,
+    num_beams,
+    top_p,
+    max_length,
+    min_length,
+    beam_search,
+    prefix,
+    postfix,
+):
+    # Check if the image folder is provided
+    if train_data_dir == '':
+        msgbox('Image folder is missing...')
+        return
+
+    # Check if the caption file extension is provided
+    if caption_file_ext == '':
+        msgbox('Please provide an extension for the caption files.')
+        return
+
+    log.info(f'Captioning files in {train_data_dir}...')
+
+    # Construct the command to run
+    run_cmd = fr'{PYTHON} "{scriptdir}/finetune/make_captions.py"'
+    run_cmd += f' --batch_size="{int(batch_size)}"'
+    run_cmd += f' --num_beams="{int(num_beams)}"'
+    run_cmd += f' --top_p="{top_p}"'
+    run_cmd += f' --max_length="{int(max_length)}"'
+    run_cmd += f' --min_length="{int(min_length)}"'
+    if beam_search:
+        run_cmd += f' --beam_search'
+    if caption_file_ext != '':
+        run_cmd += f' --caption_extension="{caption_file_ext}"'
+    run_cmd += f' "{train_data_dir}"'
+    run_cmd += f' --caption_weights="https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_large_caption.pth"'
+
+    log.info(run_cmd)
+
+    # Run the command
+    if os.name == 'posix':
+        os.system(run_cmd)
+    else:
+        subprocess.run(run_cmd)
+
+    # Add prefix and postfix
+    add_pre_postfix(
+        folder=train_data_dir,
+        caption_file_ext=caption_file_ext,
+        prefix=prefix,
+        postfix=postfix,
+    )
+
+    log.info('...captioning done')
+
+
+###
+# Gradio UI
+###
+
+
+def gradio_blip_caption_gui_tab(headless=False):
+    with gr.Tab('BLIP Captioning'):
+        gr.Markdown(
+            'This utility uses BLIP to caption files for each image in a folder.'
+        )
+        with gr.Row():
+            train_data_dir = gr.Textbox(
+                label='Image folder to caption',
+                placeholder='Directory containing the images to caption',
+                interactive=True,
+            )
+            button_train_data_dir_input = gr.Button(
+                '📂', elem_id='open_folder_small', elem_classes=["tool"], visible=(not headless)
+            )
+            button_train_data_dir_input.click(
+                get_folder_path,
+                outputs=train_data_dir,
+                show_progress=False,
+            )
+        with gr.Row():
+            caption_file_ext = gr.Textbox(
+                label='Caption file extension',
+                placeholder='Extension for caption file, e.g., .caption, .txt',
+                value='.txt',
+                interactive=True,
+            )
+
+            prefix = gr.Textbox(
+                label='Prefix to add to BLIP caption',
+                placeholder='(Optional)',
+                interactive=True,
+            )
+
+            postfix = gr.Textbox(
+                label='Postfix to add to BLIP caption',
+                placeholder='(Optional)',
+                interactive=True,
+            )
+
+            batch_size = gr.Number(
+                value=1, label='Batch size', interactive=True
+            )
+
+        with gr.Row():
+            beam_search = gr.Checkbox(
+                label='Use beam search', interactive=True, value=True
+            )
+            num_beams = gr.Number(
+                value=1, label='Number of beams', interactive=True
+            )
+            top_p = gr.Number(value=0.9, label='Top p', interactive=True)
+            max_length = gr.Number(
+                value=75, label='Max length', interactive=True
+            )
+            min_length = gr.Number(
+                value=5, label='Min length', interactive=True
+            )
+
+        caption_button = gr.Button('Caption images')
+
+        caption_button.click(
+            caption_images,
+            inputs=[
+                train_data_dir,
+                caption_file_ext,
+                batch_size,
+                num_beams,
+                top_p,
+                max_length,
+                min_length,
+                beam_search,
+                prefix,
+                postfix,
+            ],
+            show_progress=False,
+        )