extensions/sd-webui-controlnet/scripts/controlnet_ui/preset.py

import os
import gradio as gr

from typing import Dict, List

from modules import scripts
from scripts.infotext import parse_unit, serialize_unit
from scripts.controlnet_ui.tool_button import ToolButton
from scripts.logging import logger
from scripts.processor import preprocessor_filters
from scripts import external_code

save_symbol = "\U0001f4be"  # 💾
delete_symbol = "\U0001f5d1\ufe0f"  # 🗑️
refresh_symbol = "\U0001f504"  # 🔄
reset_symbol = "\U000021A9"  # ↩

NEW_PRESET = "New Preset"


def load_presets(preset_dir: str) -> Dict[str, str]:
    if not os.path.exists(preset_dir):
        os.makedirs(preset_dir)
        return {}

    presets = {}
    for filename in os.listdir(preset_dir):
        if filename.endswith(".txt"):
            with open(os.path.join(preset_dir, filename), "r") as f:
                name = filename.replace(".txt", "")
                if name == NEW_PRESET:
                    continue
                presets[name] = f.read()
    return presets


def infer_control_type(module: str, model: str) -> str:
    def matches_control_type(input_string: str, control_type: str) -> bool:
        return any(t.lower() in input_string for t in control_type.split("/"))

    control_types = preprocessor_filters.keys()
    control_type_candidates = [
        control_type
        for control_type in control_types
        if (
            matches_control_type(module, control_type)
            or matches_control_type(model, control_type)
        )
    ]
    if len(control_type_candidates) != 1:
        raise ValueError(
            f"Unable to infer control type from module {module} and model {model}"
        )
    return control_type_candidates[0]


class ControlNetPresetUI(object):
    preset_directory = os.path.join(scripts.basedir(), "presets")
    presets = load_presets(preset_directory)

    def __init__(self, id_prefix: str):
        with gr.Row():
            self.dropdown = gr.Dropdown(
                label="Presets",
                show_label=True,
                elem_classes=["cnet-preset-dropdown"],
                choices=ControlNetPresetUI.dropdown_choices(),
                value=NEW_PRESET,
            )
            self.reset_button = ToolButton(
                value=reset_symbol,
                elem_classes=["cnet-preset-reset"],
                tooltip="Reset preset",
                visible=False,
            )
            self.save_button = ToolButton(
                value=save_symbol,
                elem_classes=["cnet-preset-save"],
                tooltip="Save preset",
            )
            self.delete_button = ToolButton(
                value=delete_symbol,
                elem_classes=["cnet-preset-delete"],
                tooltip="Delete preset",
            )
            self.refresh_button = ToolButton(
                value=refresh_symbol,
                elem_classes=["cnet-preset-refresh"],
                tooltip="Refresh preset",
            )

        with gr.Box(
            elem_classes=["popup-dialog", "cnet-preset-enter-name"],
            elem_id=f"{id_prefix}_cnet_preset_enter_name",
        ) as self.name_dialog:
            with gr.Row():
                self.preset_name = gr.Textbox(
                    label="Preset name",
                    show_label=True,
                    lines=1,
                    elem_classes=["cnet-preset-name"],
                )
                self.confirm_preset_name = ToolButton(
                    value=save_symbol,
                    elem_classes=["cnet-preset-confirm-name"],
                    tooltip="Save preset",
                )

    def register_callbacks(
        self,
        uigroup,
        control_type: gr.Radio,
        *ui_states,
    ):
        def apply_preset(name: str, control_type: str, *ui_states):
            if name == NEW_PRESET:
                return (
                    gr.update(visible=False),
                    *(
                        (gr.skip(),)
                        * (len(vars(external_code.ControlNetUnit()).keys()) + 1)
                    ),
                )

            assert name in ControlNetPresetUI.presets

            infotext = ControlNetPresetUI.presets[name]
            preset_unit = parse_unit(infotext)
            current_unit = external_code.ControlNetUnit(*ui_states)
            preset_unit.image = None
            current_unit.image = None

            # Do not compare module param that are not used in preset.
            for module_param in ("processor_res", "threshold_a", "threshold_b"):
                if getattr(preset_unit, module_param) == -1:
                    setattr(current_unit, module_param, -1)

            # No update necessary.
            if vars(current_unit) == vars(preset_unit):
                return (
                    gr.update(visible=False),
                    *(
                        (gr.skip(),)
                        * (len(vars(external_code.ControlNetUnit()).keys()) + 1)
                    ),
                )

            unit = preset_unit

            try:
                new_control_type = infer_control_type(unit.module, unit.model)
            except ValueError as e:
                logger.error(e)
                new_control_type = control_type

            if new_control_type != control_type:
                uigroup.prevent_next_n_module_update += 1

            if preset_unit.module != current_unit.module:
                uigroup.prevent_next_n_slider_value_update += 1

            if preset_unit.pixel_perfect != current_unit.pixel_perfect:
                uigroup.prevent_next_n_slider_value_update += 1

            return (
                gr.update(visible=True),
                gr.update(value=new_control_type),
                *[
                    gr.update(value=value) if value is not None else gr.update()
                    for value in vars(unit).values()
                ],
            )

        for element, action in (
            (self.dropdown, "change"),
            (self.reset_button, "click"),
        ):
            getattr(element, action)(
                fn=apply_preset,
                inputs=[self.dropdown, control_type, *ui_states],
                outputs=[self.delete_button, control_type, *ui_states],
                show_progress="hidden",
            ).then(
                fn=lambda: gr.update(visible=False),
                inputs=None,
                outputs=[self.reset_button],
            )

        def save_preset(name: str, *ui_states):
            if name == NEW_PRESET:
                return gr.update(visible=True), gr.update(), gr.update()

            ControlNetPresetUI.save_preset(
                name, external_code.ControlNetUnit(*ui_states)
            )
            return (
                gr.update(),  # name dialog
                gr.update(choices=ControlNetPresetUI.dropdown_choices(), value=name),
                gr.update(visible=False),  # Reset button
            )

        self.save_button.click(
            fn=save_preset,
            inputs=[self.dropdown, *ui_states],
            outputs=[self.name_dialog, self.dropdown, self.reset_button],
            show_progress="hidden",
        ).then(
            fn=None,
            _js=f"""
            (name) => {{
                if (name === "{NEW_PRESET}")
                    popup(gradioApp().getElementById('{self.name_dialog.elem_id}'));
            }}""",
            inputs=[self.dropdown],
        )

        def delete_preset(name: str):
            ControlNetPresetUI.delete_preset(name)
            return gr.Dropdown.update(
                choices=ControlNetPresetUI.dropdown_choices(),
                value=NEW_PRESET,
            ), gr.update(visible=False)

        self.delete_button.click(
            fn=delete_preset,
            inputs=[self.dropdown],
            outputs=[self.dropdown, self.reset_button],
            show_progress="hidden",
        )

        self.name_dialog.visible = False

        def save_new_preset(new_name: str, *ui_states):
            if new_name == NEW_PRESET:
                logger.warn(f"Cannot save preset with reserved name '{NEW_PRESET}'")
                return gr.update(visible=False), gr.update()

            ControlNetPresetUI.save_preset(
                new_name, external_code.ControlNetUnit(*ui_states)
            )
            return gr.update(visible=False), gr.update(
                choices=ControlNetPresetUI.dropdown_choices(), value=new_name
            )

        self.confirm_preset_name.click(
            fn=save_new_preset,
            inputs=[self.preset_name, *ui_states],
            outputs=[self.name_dialog, self.dropdown],
            show_progress="hidden",
        ).then(fn=None, _js="closePopup")

        self.refresh_button.click(
            fn=ControlNetPresetUI.refresh_preset,
            inputs=None,
            outputs=[self.dropdown],
            show_progress="hidden",
        )

        def update_reset_button(preset_name: str, *ui_states):
            if preset_name == NEW_PRESET:
                return gr.update(visible=False)

            infotext = ControlNetPresetUI.presets[preset_name]
            preset_unit = parse_unit(infotext)
            current_unit = external_code.ControlNetUnit(*ui_states)
            preset_unit.image = None
            current_unit.image = None

            # Do not compare module param that are not used in preset.
            for module_param in ("processor_res", "threshold_a", "threshold_b"):
                if getattr(preset_unit, module_param) == -1:
                    setattr(current_unit, module_param, -1)

            return gr.update(visible=vars(current_unit) != vars(preset_unit))

        for ui_state in ui_states:
            if isinstance(ui_state, gr.Image):
                continue

            for action in ("edit", "click", "change", "clear", "release"):
                if action == "release" and not isinstance(ui_state, gr.Slider):
                    continue

                if hasattr(ui_state, action):
                    getattr(ui_state, action)(
                        fn=update_reset_button,
                        inputs=[self.dropdown, *ui_states],
                        outputs=[self.reset_button],
                    )

    @staticmethod
    def dropdown_choices() -> List[str]:
        return list(ControlNetPresetUI.presets.keys()) + [NEW_PRESET]

    @staticmethod
    def save_preset(name: str, unit: external_code.ControlNetUnit):
        infotext = serialize_unit(unit)
        with open(
            os.path.join(ControlNetPresetUI.preset_directory, f"{name}.txt"), "w"
        ) as f:
            f.write(infotext)

        ControlNetPresetUI.presets[name] = infotext

    @staticmethod
    def delete_preset(name: str):
        if name not in ControlNetPresetUI.presets:
            return

        del ControlNetPresetUI.presets[name]

        file = os.path.join(ControlNetPresetUI.preset_directory, f"{name}.txt")
        if os.path.exists(file):
            os.unlink(file)

    @staticmethod
    def refresh_preset():
        ControlNetPresetUI.presets = load_presets(ControlNetPresetUI.preset_directory)
        return gr.update(choices=ControlNetPresetUI.dropdown_choices())