kohya_gui/basic_caption_gui.py

import gradio as gr
import subprocess
from .common_gui import (
    get_folder_path,
    add_pre_postfix,
    find_replace,
    scriptdir,
    list_dirs,
    setup_environment,
)
import os
import sys

from .custom_logging import setup_logging

# Set up logging
log = setup_logging()

PYTHON = sys.executable


def caption_images(
    caption_text: str,
    images_dir: str,
    overwrite: bool,
    caption_ext: str,
    prefix: str,
    postfix: str,
    find_text: str,
    replace_text: str,
):
    """
    Captions images in a given directory with a given caption text.

    Args:
        caption_text (str): The text to be used as the caption.
        images_dir (str): The directory containing the images to be captioned.
        overwrite (bool): Whether to overwrite existing captions.
        caption_ext (str): The file extension for the caption files.
        prefix (str): Text to be added before the caption text.
        postfix (str): Text to be added after the caption text.
        find_text (str): Text to be replaced in the caption files.
        replace_text (str): Text to replace the found text in the caption files.

    Returns:
        None
    """
    # Check if images_dir and caption_ext are provided
    missing_parameters = []
    if not images_dir:
        missing_parameters.append("image directory")
    if not caption_ext:
        missing_parameters.append("caption file extension")
        
    if missing_parameters:
        log.info(
            "The following parameter(s) are missing: {}. "
            "Please provide these to proceed with captioning the images.".format(", ".join(missing_parameters))
        )
        return

    # Log the captioning process
    if caption_text:
        log.info(f"Captioning files in {images_dir} with {caption_text}...")

        # Build the command to run caption.py
        run_cmd = [
            rf"{PYTHON}",
            rf"{scriptdir}/tools/caption.py",
            "--caption_text",
            caption_text,
        ]

        # Add optional flags to the command
        if overwrite:
            run_cmd.append("--overwrite")
        if caption_ext:
            run_cmd.append("--caption_file_ext")
            run_cmd.append(caption_ext)

        run_cmd.append(rf"{images_dir}")

        # Reconstruct the safe command string for display
        command_to_run = " ".join(run_cmd)
        log.info(f"Executing command: {command_to_run}")

        # Set the environment variable for the Python path
        env = setup_environment()

        # Run the command in the sd-scripts folder context
        subprocess.run(run_cmd, env=env, shell=False)

    # Check if overwrite option is enabled
    if overwrite:
        # Add prefix and postfix to caption files or find and replace text in caption files
        if prefix or postfix or find_text:
            # Add prefix and/or postfix to caption files
            add_pre_postfix(
                folder=images_dir,
                caption_file_ext=caption_ext,
                prefix=prefix,
                postfix=postfix,
            )
            # Replace specified text in caption files if find and replace text is provided
            if find_text and replace_text:
                find_replace(
                    folder_path=images_dir,
                    caption_file_ext=caption_ext,
                    search_text=find_text,
                    replace_text=replace_text,
                )
    else:
        # Show a message if modification is not possible without overwrite option enabled
        if prefix or postfix:
            log.info(
                'Could not modify caption files with requested change because the "Overwrite existing captions in folder" option is not selected.'
            )

    # Log the end of the captioning process
    log.info("Captioning done.")


# Gradio UI
def gradio_basic_caption_gui_tab(headless=False, default_images_dir=None):
    """
    Creates a Gradio tab for basic image captioning.

    Args:
        headless (bool, optional): If True, the GUI will be headless (no visible elements). Defaults to False.
        default_images_dir (str, optional): The default directory to use for image selection. If not provided,
            it defaults to the 'data' directory in the script directory.

    Returns:
        None
    """
    from .common_gui import create_refresh_button

    # Set default images directory if not provided
    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

    # Function to list directories
    def list_images_dirs(path):
        """
        Lists directories within a specified path and updates the current image directory.

        Parameters:
            path (str): The directory path to list image directories from.

        Returns:
            list: A list of directories within the specified 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))

    # Gradio tab for basic captioning
    with gr.Tab("Basic Captioning"):
        # Markdown description
        gr.Markdown(
            "This utility allows you to create simple caption files for each image in a folder."
        )
        # Group and row for image folder selection
        with gr.Group(), gr.Row():
            # Dropdown for image folder
            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,
            )
            # Refresh button for image folder
            create_refresh_button(
                images_dir,
                lambda: None,
                lambda: {"choices": list_images_dirs(current_images_dir)},
                "open_folder_small",
            )
            # Button to open folder
            folder_button = gr.Button(
                "📂",
                elem_id="open_folder_small",
                elem_classes=["tool"],
                visible=(not headless),
            )
            # Event handler for button click
            folder_button.click(
                get_folder_path,
                outputs=images_dir,
                show_progress=False,
            )
            # Textbox for caption file extension
            caption_ext = gr.Dropdown(
                label="Caption file extension",
                choices=[".cap", ".caption", ".txt"],
                value=".txt",
                interactive=True,
                allow_custom_value=True,
            )
            # Checkbox to overwrite existing captions
            overwrite = gr.Checkbox(
                label="Overwrite existing captions in folder",
                interactive=True,
                value=False,
            )
        # Row for caption prefix and text
        with gr.Row():
            # Textbox for caption prefix
            prefix = gr.Textbox(
                label="Prefix to add to caption",
                placeholder="(Optional)",
                interactive=True,
            )
            # Textbox for caption text
            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,
                lines=2,
            )
            # Textbox for caption postfix
            postfix = gr.Textbox(
                label="Postfix to add to caption",
                placeholder="(Optional)",
                interactive=True,
            )
        # Group and row for find and replace text
        with gr.Group(), gr.Row():
            # Textbox for find text
            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,
                lines=2,
            )
            # Textbox for replace text
            replace_text = gr.Textbox(
                label="Replacement text",
                placeholder='e.g., "by some artist". Leave empty if you want to replace with nothing.',
                interactive=True,
                lines=2,
            )
            # Button to caption images
            caption_button = gr.Button("Caption images")
            # Event handler for button click
            caption_button.click(
                caption_images,
                inputs=[
                    caption_text,
                    images_dir,
                    overwrite,
                    caption_ext,
                    prefix,
                    postfix,
                    find_text,
                    replace_text,
                ],
                show_progress=False,
            )

        # Event handler for dynamic update of dropdown choices
        images_dir.change(
            fn=lambda path: gr.Dropdown(choices=[""] + list_images_dirs(path)),
            inputs=images_dir,
            outputs=images_dir,
            show_progress=False,
        )