kohya_gui/class_source_model.py

import gradio as gr
import os

from .common_gui import (
    get_file_path,
    get_folder_path,
    set_pretrained_model_name_or_path_input,
    scriptdir,
    list_dirs,
    list_files,
    create_refresh_button,
)
from .class_gui_config import KohyaSSGUIConfig

folder_symbol = "\U0001f4c2"  # 📂
refresh_symbol = "\U0001f504"  # 🔄
save_style_symbol = "\U0001f4be"  # 💾
document_symbol = "\U0001F4C4"  # 📄

default_models = [
    "stabilityai/stable-diffusion-xl-base-1.0",
    "stabilityai/stable-diffusion-xl-refiner-1.0",
    "stabilityai/stable-diffusion-2-1-base/blob/main/v2-1_512-ema-pruned",
    "stabilityai/stable-diffusion-2-1-base",
    "stabilityai/stable-diffusion-2-base",
    "stabilityai/stable-diffusion-2-1/blob/main/v2-1_768-ema-pruned",
    "stabilityai/stable-diffusion-2-1",
    "stabilityai/stable-diffusion-2",
    "runwayml/stable-diffusion-v1-5",
    "CompVis/stable-diffusion-v1-4",
]


class SourceModel:
    def __init__(
        self,
        save_model_as_choices=[
            "same as source model",
            "ckpt",
            "diffusers",
            "diffusers_safetensors",
            "safetensors",
        ],
        save_precision_choices=[
            "float",
            "fp16",
            "bf16",
        ],
        headless=False,
        finetuning=False,
        config: KohyaSSGUIConfig = {},
    ):
        self.headless = headless
        self.save_model_as_choices = save_model_as_choices
        self.finetuning = finetuning
        self.config = config

        # Set default directories if not provided
        self.current_models_dir = self.config.get(
            "model.models_dir", os.path.join(scriptdir, "models")
        )
        self.current_train_data_dir = self.config.get(
            "model.train_data_dir", os.path.join(scriptdir, "data")
        )
        self.current_dataset_config_dir = self.config.get(
            "model.dataset_config", os.path.join(scriptdir, "dataset_config")
        )

        model_checkpoints = list(
            list_files(
                self.current_models_dir, exts=[".ckpt", ".safetensors"], all=True
            )
        )

        def list_models(path):
            self.current_models_dir = (
                path if os.path.isdir(path) else os.path.dirname(path)
            )
            return default_models + list(
                list_files(path, exts=[".ckpt", ".safetensors"], all=True)
            )

        def list_train_data_dirs(path):
            self.current_train_data_dir = path if not path == "" else "."
            return list(list_dirs(self.current_train_data_dir))

        def list_dataset_config_dirs(path: str) -> list:
            """
            List directories and toml files in the dataset_config directory.

            Parameters:
            - path (str): The path to list directories and files from.

            Returns:
            - list: A list of directories and files.
            """
            current_dataset_config_dir = path if not path == "" else "."
            # Lists all .json files in the current configuration directory, used for populating dropdown choices.
            return list(
                list_files(current_dataset_config_dir, exts=[".toml"], all=True)
            )

        with gr.Accordion("Model", open=True):
            with gr.Column(), gr.Group():
                model_ext = gr.Textbox(value="*.safetensors *.ckpt", visible=False)
                model_ext_name = gr.Textbox(value="Model types", visible=False)

                # Define the input elements
                with gr.Row():
                    with gr.Column(), gr.Row():
                        self.model_list = gr.Textbox(visible=False, value="")
                        self.pretrained_model_name_or_path = gr.Dropdown(
                            label="Pretrained model name or path",
                            choices=default_models + model_checkpoints,
                            value=self.config.get("model.models_dir", "runwayml/stable-diffusion-v1-5"),
                            allow_custom_value=True,
                            visible=True,
                            min_width=100,
                        )
                        create_refresh_button(
                            self.pretrained_model_name_or_path,
                            lambda: None,
                            lambda: {"choices": list_models(self.current_models_dir)},
                            "open_folder_small",
                        )

                        self.pretrained_model_name_or_path_file = gr.Button(
                            document_symbol,
                            elem_id="open_folder_small",
                            elem_classes=["tool"],
                            visible=(not headless),
                        )
                        self.pretrained_model_name_or_path_file.click(
                            get_file_path,
                            inputs=[self.pretrained_model_name_or_path, model_ext, model_ext_name],
                            outputs=self.pretrained_model_name_or_path,
                            show_progress=False,
                        )
                        self.pretrained_model_name_or_path_folder = gr.Button(
                            folder_symbol,
                            elem_id="open_folder_small",
                            elem_classes=["tool"],
                            visible=(not headless),
                        )
                        self.pretrained_model_name_or_path_folder.click(
                            get_folder_path,
                            inputs=self.pretrained_model_name_or_path,
                            outputs=self.pretrained_model_name_or_path,
                            show_progress=False,
                        )

                    with gr.Column(), gr.Row():
                        self.output_name = gr.Textbox(
                            label="Trained Model output name",
                            placeholder="(Name of the model to output)",
                            value=self.config.get("model.output_name", "last"),
                            interactive=True,
                        )
                with gr.Row():
                    with gr.Column(), gr.Row():
                        self.train_data_dir = gr.Dropdown(
                            label=(
                                "Image folder (containing training images subfolders)"
                                if not finetuning
                                else "Image folder (containing training images)"
                            ),
                            choices=[""]
                            + list_train_data_dirs(self.current_train_data_dir),
                            value=self.config.get("model.train_data_dir", ""),
                            interactive=True,
                            allow_custom_value=True,
                        )
                        create_refresh_button(
                            self.train_data_dir,
                            lambda: None,
                            lambda: {
                                "choices": [""]
                                + list_train_data_dirs(self.current_train_data_dir)
                            },
                            "open_folder_small",
                        )
                        self.train_data_dir_folder = gr.Button(
                            "📂",
                            elem_id="open_folder_small",
                            elem_classes=["tool"],
                            visible=(not self.headless),
                        )
                        self.train_data_dir_folder.click(
                            get_folder_path,
                            outputs=self.train_data_dir,
                            show_progress=False,
                        )
                    with gr.Column(), gr.Row():
                        # Toml directory dropdown
                        self.dataset_config = gr.Dropdown(
                            label="Dataset config file (Optional. Select the toml configuration file to use for the dataset)",
                            choices=[self.config.get("model.dataset_config", "")]
                            + list_dataset_config_dirs(self.current_dataset_config_dir),
                            value=self.config.get("model.dataset_config", ""),
                            interactive=True,
                            allow_custom_value=True,
                        )
                        # Refresh button for dataset_config directory
                        create_refresh_button(
                            self.dataset_config,
                            lambda: None,
                            lambda: {
                                "choices": [""]
                                + list_dataset_config_dirs(
                                    self.current_dataset_config_dir
                                )
                            },
                            "open_folder_small",
                        )
                        # Toml directory button
                        self.dataset_config_folder = gr.Button(
                            document_symbol,
                            elem_id="open_folder_small",
                            elem_classes=["tool"],
                            visible=(not self.headless),
                        )

                        # Toml directory button click event
                        self.dataset_config_folder.click(
                            get_file_path,
                            inputs=[
                                self.dataset_config,
                                gr.Textbox(value="*.toml", visible=False),
                                gr.Textbox(value="Dataset config types", visible=False),
                            ],
                            outputs=self.dataset_config,
                            show_progress=False,
                        )
                        # Change event for dataset_config directory dropdown
                        self.dataset_config.change(
                            fn=lambda path: gr.Dropdown(
                                choices=[""] + list_dataset_config_dirs(path)
                            ),
                            inputs=self.dataset_config,
                            outputs=self.dataset_config,
                            show_progress=False,
                        )

                with gr.Row():
                    with gr.Column():
                        with gr.Row():
                            self.v2 = gr.Checkbox(
                                label="v2", value=False, visible=False, min_width=60
                            )
                            self.v_parameterization = gr.Checkbox(
                                label="v_parameterization",
                                value=False,
                                visible=False,
                                min_width=130,
                            )
                            self.sdxl_checkbox = gr.Checkbox(
                                label="SDXL",
                                value=False,
                                visible=False,
                                min_width=60,
                            )
                    with gr.Column():
                        gr.Group(visible=False)

                with gr.Row():
                    self.training_comment = gr.Textbox(
                        label="Training comment",
                        placeholder="(Optional) Add training comment to be included in metadata",
                        interactive=True,
                        value=self.config.get("model.training_comment", ""),
                    )

                with gr.Row():
                    self.save_model_as = gr.Radio(
                        save_model_as_choices,
                        label="Save trained model as",
                        value=self.config.get("model.save_model_as", "safetensors"),
                    )
                    self.save_precision = gr.Radio(
                        save_precision_choices,
                        label="Save precision",
                        value=self.config.get("model.save_precision", "fp16"),
                    )

                self.pretrained_model_name_or_path.change(
                    fn=lambda path: set_pretrained_model_name_or_path_input(
                        path, refresh_method=list_models
                    ),
                    inputs=[
                        self.pretrained_model_name_or_path,
                    ],
                    outputs=[
                        self.pretrained_model_name_or_path,
                        self.v2,
                        self.v_parameterization,
                        self.sdxl_checkbox,
                    ],
                    show_progress=False,
                )

                self.train_data_dir.change(
                    fn=lambda path: gr.Dropdown(
                        choices=[""] + list_train_data_dirs(path)
                    ),
                    inputs=self.train_data_dir,
                    outputs=self.train_data_dir,
                    show_progress=False,
                )