a1111 implementation of https://github.com/ChenyangSi/FreeU
import dataclasses
import inspect
import json
import pathlib
import re
import sys
from typing import Union, List, Any


@dataclasses.dataclass
class StageInfo:
    backbone_factor: float = 1.0
    skip_factor: float = 1.0
    backbone_offset: float = 0.0
    backbone_width: float = 0.5
    skip_cutoff: float = 0.0
    skip_high_end_factor: float = 1.0
    # <- add new fields at the end here for png info backwards compatibility

    def to_dict(self, include_default=False):
        default_stage_info = StageInfo()
        res = vars(self).copy()
        for k, v in res.copy().items():
            if not include_default and v == getattr(default_stage_info, k):
                del res[k]

        return res

    def copy(self):
        return StageInfo(**vars(self))


STAGE_INFO_ARGS_LEN = len(inspect.getfullargspec(StageInfo.__init__)[0]) - 1  # off by one because of self
STAGES_COUNT = 3
shorthand_re = re.compile(r"^([a-z]{1,2})([0-9]+)$")

xyz_attrs: dict = {}
current_sampling_step: float = 0


@dataclasses.dataclass
class State:
    enable: bool = True
    start_ratio: Union[float, int] = 0.0
    stop_ratio: Union[float, int] = 1.0
    transition_smoothness: float = 0.0
    stage_infos: List[Union[StageInfo, dict, Any]] = dataclasses.field(default_factory=lambda: [StageInfo() for _ in range(STAGES_COUNT)])

    def __post_init__(self):
        self.stage_infos = self.group_stage_infos()

    def group_stage_infos(self):
        res = []
        i = 0
        while i < len(self.stage_infos) and len(res) < STAGES_COUNT:
            if isinstance(self.stage_infos[i], StageInfo):
                res.append(self.stage_infos[i])
                i += 1
            elif isinstance(self.stage_infos[i], dict):
                res.append(StageInfo(**self.stage_infos[i]))
                i += 1
            else:
                next_i = i + STAGE_INFO_ARGS_LEN
                res.append(StageInfo(*self.stage_infos[i:next_i]))
                i = next_i

        for _ in range(STAGES_COUNT - len(res)):
            res.append(StageInfo())

        return res

    def to_dict(self):
        result = vars(self).copy()
        result["stage_infos"] = [stage_info.to_dict() for stage_info in result["stage_infos"]]
        del result["enable"]
        return result

    def copy(self):
        self_vars = vars(self)
        old_stage_infos = self_vars["stage_infos"]
        self_vars["stage_infos"] = old_stage_infos.copy()
        for i, stage_info in enumerate(old_stage_infos):
            self_vars["stage_infos"][i] = stage_info.copy()

        return State(**self_vars)

    def update_attr(self, key, value):
        if match := shorthand_re.match(key):
            char, index = match.group(1, 2)
            stage_info = self.stage_infos[int(index)]
            if char == "b":
                stage_info.backbone_factor = value
            elif char == "s":
                stage_info.skip_factor = value
            elif char == "o":
                stage_info.backbone_offset = value
            elif char == "w":
                stage_info.backbone_width = value
            elif char == "t":
                stage_info.skip_cutoff = value
            elif char == "h":
                stage_info.skip_high_end_factor = value
        else:
            self.__dict__[key] = value


def apply_xyz():
    global instance

    if preset_key := xyz_attrs.get("preset"):
        if preset := all_presets.get(preset_key):
            instance = preset.copy()
        elif preset_key != "UI Settings":
            print("[sd-webui-freeu]", f"XYZ Preset '{preset_key}' does not exist", file=sys.stderr)

    for k, v in xyz_attrs.items():
        if k == "preset":
            continue

        instance.update_attr(k, v)


STATE_ARGS_LEN = len(inspect.getfullargspec(State.__init__)[0]) - 1  # off by one because of self
PRESETS_PATH = pathlib.Path(__file__).parent.parent / "presets.json"

instance = State()
default_presets = {
    "SD1.4 Recommendations": State(
        stage_infos=[
            StageInfo(1.2, 0.9),
            StageInfo(1.4, 0.2),
            StageInfo(1, 1),
        ],
    ),
    "SD2.1 Recommendations": State(
        stage_infos=[
            StageInfo(1.1, 0.9),
            StageInfo(1.2, 0.2),
            StageInfo(1, 1),
        ],
    ),
    "SDXL Recommendations": State(
        stage_infos=[
            StageInfo(1.1, 0.6),
            StageInfo(1.2, 0.4),
            StageInfo(1, 1),
        ],
    ),
}
all_presets = {}


def reload_presets():
    all_presets.clear()
    all_presets.update(default_presets)
    all_presets.update(load_presets())


def load_presets():
    if not PRESETS_PATH.exists():
        return []

    with open(PRESETS_PATH, "r") as f:
        return {
            k: State(**v)
            for k, v in json.load(f).items()
        }


def save_presets(presets=None):
    if presets is None:
        presets = get_user_presets()

    presets = {k: v.to_dict() for k, v in presets.items()}

    with open(PRESETS_PATH, "w") as f:
        json.dump(presets, f)


def get_user_presets():
    return {
        k: v
        for k, v in all_presets.items()
        if k not in default_presets
    }

lib_free_u/global_state.py
ljleb/sd-webui-freeu
import functools
import math
import pathlib
import sys
from typing import Tuple, Union
from lib_free_u import global_state
from modules import scripts, shared
from modules.sd_hijack_unet import th
import torch


def patch():
    th.cat = functools.partial(free_u_cat_hijack, original_function=th.cat)

    cn_script_paths = [
        str(pathlib.Path(scripts.basedir()).parent.parent / "extensions-builtin" / "sd-webui-controlnet"),
        str(pathlib.Path(scripts.basedir()).parent / "sd-webui-controlnet"),
    ]
    sys.path[0:0] = cn_script_paths
    cn_status = "enabled"
    try:
        import scripts.hook as controlnet_hook
    except ImportError:
        cn_status = "disabled"
    else:
        controlnet_hook.th.cat = functools.partial(free_u_cat_hijack, original_function=controlnet_hook.th.cat)
    finally:
        for p in cn_script_paths:
            sys.path.remove(p)

        print("[sd-webui-freeu]", f"Controlnet support: *{cn_status}*")


def free_u_cat_hijack(hs, *args, original_function, **kwargs):
    if not global_state.instance.enable:
        return original_function(hs, *args, **kwargs)

    schedule_ratio = get_schedule_ratio()
    if schedule_ratio == 0:
        return original_function(hs, *args, **kwargs)

    try:
        h, h_skip = hs
        if list(kwargs.keys()) != ["dim"] or kwargs.get("dim", -1) != 1:
            return original_function(hs, *args, **kwargs)
    except ValueError:
        return original_function(hs, *args, **kwargs)

    dims = h.shape[1]
    try:
        index = [1280, 640, 320].index(dims)
        stage_info = global_state.instance.stage_infos[index]
    except ValueError:
        stage_info = None

    if stage_info is not None:
        redion_begin, region_end, region_inverted = ratio_to_region(stage_info.backbone_width, stage_info.backbone_offset, dims)
        mask = torch.arange(dims)
        mask = (redion_begin <= mask) & (mask <= region_end)
        if region_inverted:
            mask = ~mask

        h[:, mask] *= lerp(1, stage_info.backbone_factor, schedule_ratio)
        h_skip = filter_skip(
            h_skip,
            threshold=stage_info.skip_cutoff,
            scale=lerp(1, stage_info.skip_factor, schedule_ratio),
            scale_high=lerp(1, stage_info.skip_high_end_factor, schedule_ratio),
        )

    return original_function([h, h_skip], *args, **kwargs)


def filter_skip(x, threshold, scale, scale_high):
    if scale == 1 and scale_high == 1:
        return x

    fft_device = x.device
    if no_gpu_complex_support():
        fft_device = "cpu"

    # FFT
    x_freq = torch.fft.fftn(x.to(fft_device).float(), dim=(-2, -1))
    x_freq = torch.fft.fftshift(x_freq, dim=(-2, -1))

    B, C, H, W = x_freq.shape
    mask = torch.full((B, C, H, W), float(scale_high), device=fft_device)

    crow, ccol = H // 2, W // 2
    threshold_row = max(1, math.floor(crow * threshold))
    threshold_col = max(1, math.floor(ccol * threshold))
    mask[..., crow - threshold_row:crow + threshold_row, ccol - threshold_col:ccol + threshold_col] = scale
    x_freq *= mask

    # IFFT
    x_freq = torch.fft.ifftshift(x_freq, dim=(-2, -1))
    x_filtered = torch.fft.ifftn(x_freq, dim=(-2, -1)).real.to(device=x.device, dtype=x.dtype)

    return x_filtered


def ratio_to_region(width: float, offset: float, n: int) -> Tuple[int, int, bool]:
    if width < 0:
        offset += width
        width = -width
    width = min(width, 1)

    if offset < 0:
        offset = 1 + offset - int(offset)
    offset = math.fmod(offset, 1.0)

    if width + offset <= 1:
        inverted = False
        start = offset * n
        end = (width + offset) * n
    else:
        inverted = True
        start = (width + offset - 1) * n
        end = offset * n

    return round(start), round(end), inverted


def get_schedule_ratio():
    start_step = to_denoising_step(global_state.instance.start_ratio)
    stop_step = to_denoising_step(global_state.instance.stop_ratio)

    if start_step == stop_step:
        smooth_schedule_ratio = 0.0
    elif global_state.current_sampling_step < start_step:
        smooth_schedule_ratio = min(1.0, max(0.0, global_state.current_sampling_step / start_step))
    else:
        smooth_schedule_ratio = min(1.0, max(0.0, 1 + (global_state.current_sampling_step - start_step) / (start_step - stop_step)))

    flat_schedule_ratio = 1.0 if start_step <= global_state.current_sampling_step < stop_step else 0.0

    return lerp(flat_schedule_ratio, smooth_schedule_ratio, global_state.instance.transition_smoothness)


def to_denoising_step(number: Union[float, int], steps=None) -> int:
    if steps is None:
        steps = shared.state.sampling_steps

    if isinstance(number, float):
        return int(number * steps)

    return number


def lerp(a, b, r):
    return (1-r)*a + r*b


def no_gpu_complex_support():
    mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
    try:
        import torch_directml
    except ImportError:
        dml_available = False
    else:
        dml_available = torch_directml.is_available()

    return mps_available or dml_available

lib_free_u/unet.py
ljleb/sd-webui-freeu
import sys
from types import ModuleType
from typing import Optional
from modules import scripts
from lib_free_u import global_state


def patch():
    xyz_module = find_xyz_module()
    if xyz_module is None:
        print("[sd-webui-freeu]", "xyz_grid.py not found.", file=sys.stderr)
        return

    xyz_module.axis_options.extend([
        xyz_module.AxisOption("[FreeU] Enabled", str_to_bool, apply_global_state("enable"), choices=choices_bool),
        xyz_module.AxisOption("[FreeU] Preset", str, apply_global_state("preset"), choices=choices_preset),
        xyz_module.AxisOption("[FreeU] Start At Step", int_or_float, apply_global_state("start_ratio")),
        xyz_module.AxisOption("[FreeU] Stop At Step", int_or_float, apply_global_state("stop_ratio")),
        xyz_module.AxisOption("[FreeU] Transition Smoothness", int_or_float, apply_global_state("transition_smoothness")),
        *[
            opt
            for index in range(global_state.STAGES_COUNT)
            for opt in [
                xyz_module.AxisOption(f"[FreeU] Stage {index+1} Backbone Scale", float, apply_global_state(f"b{index}")),
                xyz_module.AxisOption(f"[FreeU] Stage {index+1} Backbone Offset", float, apply_global_state(f"o{index}")),
                xyz_module.AxisOption(f"[FreeU] Stage {index+1} Backbone Width", float, apply_global_state(f"w{index}")),
                xyz_module.AxisOption(f"[FreeU] Stage {index+1} Skip Scale", float, apply_global_state(f"s{index}")),
                xyz_module.AxisOption(f"[FreeU] Stage {index+1} Skip Cutoff", float, apply_global_state(f"t{index}")),
                xyz_module.AxisOption(f"[FreeU] Stage {index+1} Skip High End Scale", float, apply_global_state(f"h{index}")),
            ]
        ]
    ])


def apply_global_state(k):
    def callback(_p, v, _vs):
        global_state.xyz_attrs[k] = v

    return callback


def str_to_bool(string):
    string = str(string)
    if string in ["None", ""]:
        return None
    elif string.lower() in ["true", "1"]:
        return True
    elif string.lower() in ["false", "0"]:
        return False
    else:
        raise ValueError(f"Could not convert string to boolean: {string}")


def int_or_float(string):
    try:
        return int(string)
    except ValueError:
        return float(string)


def choices_bool():
    return ["False", "True"]


def choices_preset():
    presets = list(global_state.all_presets.keys())
    presets.insert(0, "UI Settings")
    return presets
    

def find_xyz_module() -> Optional[ModuleType]:
    for data in scripts.scripts_data:
        if data.script_class.__module__ in {"xyz_grid.py", "xy_grid.py"} and hasattr(data, "module"):
            return data.module

    return None

lib_free_u/xyz_grid.py
ljleb/sd-webui-freeu
import json
import gradio as gr
from modules import scripts, script_callbacks, processing, shared
from lib_free_u import global_state, unet, xyz_grid


txt2img_steps_component = None
img2img_steps_component = None
txt2img_steps_callbacks = []
img2img_steps_callbacks = []


class FreeUScript(scripts.Script):
    def title(self):
        return "FreeU"

    def show(self, is_img2img):
        return scripts.AlwaysVisible

    def ui(self, is_img2img):
        global_state.reload_presets()
        default_stage_infos = next(iter(global_state.all_presets.values())).stage_infos

        with gr.Accordion(open=False, label=self.title()):
            with gr.Row():
                enabled = gr.Checkbox(
                    label="Enable",
                    value=False,
                )
                preset_name = gr.Dropdown(
                    show_label=False,
                    choices=list(global_state.all_presets.keys()),
                    value=next(iter(global_state.all_presets.keys())),
                    type="value",
                    elem_id=self.elem_id("preset_name"),
                    allow_custom_value=True,
                    tooltip="Apply button loads settings\nWrite custom name to enable save\nDelete automatically will save to file",
                    size="sm",
                )

                is_custom_preset = preset_name.value not in global_state.default_presets
                preset_exists = preset_name.value in global_state.all_presets

                apply_preset = gr.Button(
                    value="✅",
                    size="lg",
                    elem_classes="tool",
                    interactive=preset_exists,
                )
                save_preset = gr.Button(
                    value="💾",
                    size="lg",
                    elem_classes="tool",
                    interactive=is_custom_preset,
                )
                refresh_presets = gr.Button(
                    value="🔄",
                    size="lg",
                    elem_classes="tool"
                )
                delete_preset = gr.Button(
                    value="🗑️",
                    size="lg",
                    elem_classes="tool",
                    interactive=is_custom_preset and preset_exists,
                )

            with gr.Row():
                start_ratio = gr.Slider(
                    label="Start At Step",
                    elem_id=self.elem_id("start_at_step"),
                    minimum=0,
                    maximum=1,
                    value=0,
                )

                stop_ratio = gr.Slider(
                    label="Stop At Step",
                    elem_id=self.elem_id("stop_at_step"),
                    minimum=0,
                    maximum=1,
                    value=1,
                )

                transition_smoothness = gr.Slider(
                    label="Transition Smoothness",
                    elem_id=self.elem_id("transition_smoothness"),
                    minimum=0,
                    maximum=1,
                    value=0,
                )

            flat_stage_infos = []

            for index in range(global_state.STAGES_COUNT):
                stage_n = index + 1
                default_stage_info = default_stage_infos[index]

                with gr.Accordion(open=index < 2, label=f"Stage {stage_n}"):
                    with gr.Row():
                        backbone_scale = gr.Slider(
                            label=f"Backbone {stage_n} Scale",
                            elem_id=self.elem_id(f"backbone_scale_{stage_n}"),
                            minimum=-1,
                            maximum=3,
                            value=default_stage_info.backbone_factor,
                        )

                        backbone_offset = gr.Slider(
                            label=f"Backbone {stage_n} Offset",
                            elem_id=self.elem_id(f"backbone_offset_{stage_n}"),
                            minimum=0,
                            maximum=1,
                            value=default_stage_info.backbone_offset,
                        )

                        backbone_width = gr.Slider(
                            label=f"Backbone {stage_n} Width",
                            elem_id=self.elem_id(f"backbone_width_{stage_n}"),
                            minimum=0,
                            maximum=1,
                            value=default_stage_info.backbone_width,
                        )

                    with gr.Row():
                        skip_scale = gr.Slider(
                            label=f"Skip {stage_n} Scale",
                            elem_id=self.elem_id(f"skip_scale_{stage_n}"),
                            minimum=-1,
                            maximum=3,
                            value=default_stage_info.skip_factor,
                        )

                        skip_high_end_scale = gr.Slider(
                            label=f"Skip {stage_n} High End Scale",
                            elem_id=self.elem_id(f"skip_high_end_scale_{stage_n}"),
                            minimum=-1,
                            maximum=3,
                            value=default_stage_info.skip_high_end_factor,
                        )

                        skip_cutoff = gr.Slider(
                            label=f"Skip {stage_n} Cutoff",
                            elem_id=self.elem_id(f"skip_cutoff_{stage_n}"),
                            minimum=0.0,
                            maximum=1.0,
                            value=default_stage_info.skip_cutoff,
                        )

                flat_stage_infos.extend([
                    backbone_scale,
                    skip_scale,
                    backbone_offset,
                    backbone_width,
                    skip_cutoff,
                    skip_high_end_scale,
                ])

        def on_preset_name_change(preset_name):
            is_custom_preset = preset_name not in global_state.default_presets
            preset_exists = preset_name in global_state.all_presets
            return (
                gr.Button.update(interactive=preset_exists),
                gr.Button.update(interactive=is_custom_preset),
                gr.Button.update(interactive=is_custom_preset and preset_exists),
            )

        preset_name.change(
            fn=on_preset_name_change,
            inputs=[preset_name],
            outputs=[apply_preset, save_preset, delete_preset],
        )

        def on_apply_click(user_settings_name):
            preset = global_state.all_presets[user_settings_name]
            return (
                gr.Slider.update(value=preset.start_ratio),
                gr.Slider.update(value=preset.stop_ratio),
                gr.Slider.update(value=preset.transition_smoothness),
                *[
                    gr.update(value=v)
                    for stage_info in preset.stage_infos
                    for v in stage_info.to_dict(include_default=True).values()
                ],
            )

        apply_preset.click(
            fn=on_apply_click,
            inputs=[preset_name],
            outputs=[start_ratio, stop_ratio, transition_smoothness, *flat_stage_infos],
        )

        def on_save_click(preset_name, start_ratio, stop_ratio, transition_smoothness, *flat_stage_infos):
            global_state.all_presets[preset_name] = global_state.State(
                stage_infos=flat_stage_infos,
                start_ratio=start_ratio,
                stop_ratio=stop_ratio,
                transition_smoothness=transition_smoothness,
            )
            global_state.save_presets()

            return (
                gr.Dropdown.update(choices=list(global_state.all_presets.keys())),
                gr.Button.update(interactive=True),
                gr.Button.update(interactive=True),
            )

        save_preset.click(
            fn=on_save_click,
            inputs=[preset_name, start_ratio, stop_ratio, transition_smoothness, *flat_stage_infos],
            outputs=[preset_name, apply_preset, delete_preset],
        )

        def on_refresh_click(preset_name):
            global_state.reload_presets()
            is_custom_preset = preset_name not in global_state.default_presets
            preset_exists = preset_name in global_state.all_presets

            return (
                gr.Dropdown.update(value=preset_name, choices=list(global_state.all_presets.keys())),
                gr.Button.update(interactive=preset_exists),
                gr.Button.update(interactive=is_custom_preset),
                gr.Button.update(interactive=is_custom_preset and preset_exists),
            )

        refresh_presets.click(
            fn=on_refresh_click,
            inputs=[preset_name],
            outputs=[preset_name, apply_preset, save_preset, delete_preset],
        )

        def on_delete_click(preset_name):
            preset_name_index = list(global_state.all_presets.keys()).index(preset_name)
            del global_state.all_presets[preset_name]
            global_state.save_presets()

            preset_name_index = min(len(global_state.all_presets) - 1, preset_name_index)
            preset_names = list(global_state.all_presets.keys())
            preset_name = preset_names[preset_name_index]

            is_custom_preset = preset_name not in global_state.default_presets
            preset_exists = preset_name in global_state.all_presets
            return (
                gr.Dropdown.update(value=preset_name, choices=preset_names),
                gr.Button.update(interactive=preset_exists),
                gr.Button.update(interactive=is_custom_preset),
                gr.Button.update(interactive=is_custom_preset and preset_exists),
            )

        delete_preset.click(
            fn=on_delete_click,
            inputs=[preset_name],
            outputs=[preset_name, apply_preset, save_preset, delete_preset],
        )

        schedule_infotext = gr.HTML(visible=False, interactive=False)
        stages_infotext = gr.HTML(visible=False, interactive=False)

        def register_schedule_infotext_change(steps_component):
            schedule_infotext.change(
                fn=self.on_schedule_infotext_update,
                inputs=[schedule_infotext, steps_component],
                outputs=[schedule_infotext, start_ratio, stop_ratio, transition_smoothness],
            )

        steps_component, steps_callbacks = (
            (img2img_steps_component, img2img_steps_callbacks)
            if is_img2img else
            (txt2img_steps_component, txt2img_steps_callbacks)
        )

        if steps_component is None:
            steps_callbacks.append(register_schedule_infotext_change)
        else:
            register_schedule_infotext_change(steps_component)

        stages_infotext.change(
            fn=self.on_stages_infotext_update,
            inputs=[stages_infotext],
            outputs=[stages_infotext, enabled, *flat_stage_infos],
        )

        self.infotext_fields = [
            (schedule_infotext, "FreeU Schedule"),
            (stages_infotext, "FreeU Stages"),
        ]
        self.paste_field_names = [f for _, f in self.infotext_fields]

        return enabled, start_ratio, stop_ratio, transition_smoothness, *flat_stage_infos

    def on_schedule_infotext_update(self, infotext, steps):
        if not infotext:
            return (gr.skip(),) * 4

        start_ratio, stop_ratio, transition_smoothness, *_ = infotext.split(", ")

        return (
            gr.update(value=""),
            gr.update(value=unet.to_denoising_step(xyz_grid.int_or_float(start_ratio), steps) / steps),
            gr.update(value=unet.to_denoising_step(xyz_grid.int_or_float(stop_ratio), steps) / steps),
            gr.update(value=float(transition_smoothness)),
        )

    def on_stages_infotext_update(self, infotext):
        if not infotext:
            return (gr.skip(),) * (2 + global_state.STAGES_COUNT * global_state.STAGE_INFO_ARGS_LEN)

        stage_infos = json.loads(infotext)
        stage_infos = [
            global_state.StageInfo(**stage_info)
            for stage_info in stage_infos
        ]
        stage_infos.extend([
            global_state.StageInfo()
            for _ in range(global_state.STAGES_COUNT - len(stage_infos))
        ])

        return (
            gr.update(value=""),
            gr.update(value=shared.opts.data.get("freeu_png_info_auto_enable", True)),
            *(
                gr.update(value=v)
                for stage_info in stage_infos
                for v in stage_info.to_dict(include_default=True).values()
            )
        )

    def process(
        self,
        p: processing.StableDiffusionProcessing,
        *args
    ):
        if isinstance(args[0], dict):
            global_state.instance = global_state.State(**args[0])
        elif isinstance(args[0], bool):
            stage_infos_begin = global_state.STATE_ARGS_LEN - 1
            global_state.instance = global_state.State(
                args[0],
                *[float(n) for n in args[1:stage_infos_begin]],
                args[stage_infos_begin:],
            )
        else:
            raise TypeError(f"Unrecognized args sequence starting with type {type(args[0])}")

        global_state.apply_xyz()
        global_state.xyz_attrs.clear()
        if not global_state.instance.enable:
            return

        last_d = False
        p.extra_generation_params["FreeU Stages"] = json.dumps(list(reversed([
            stage_info.to_dict()
            for stage_info in reversed(global_state.instance.stage_infos)
            # strip all empty dicts
            if last_d or stage_info.to_dict() and (last_d := True)
        ])))
        p.extra_generation_params["FreeU Schedule"] = ", ".join([
            str(global_state.instance.start_ratio),
            str(global_state.instance.stop_ratio),
            str(global_state.instance.transition_smoothness),
        ])

    def process_batch(self, p, *args, **kwargs):
        global_state.current_sampling_step = 0


def increment_sampling_step(*_args, **_kwargs):
    global_state.current_sampling_step += 1


try:
    script_callbacks.on_cfg_after_cfg(increment_sampling_step)
except AttributeError:
    # webui < 1.6.0
    # normally we should increment the current sampling step after cfg
    # but as long as we don't need to run code during cfg it should be fine to increment early
    script_callbacks.on_cfg_denoised(increment_sampling_step)


def on_after_component(component, **kwargs):
    global txt2img_steps_component, img2img_steps_component

    if kwargs.get("elem_id", None) == "img2img_steps":
        img2img_steps_component = component
        for callback in img2img_steps_callbacks:
            callback(component)

    if kwargs.get("elem_id", None) == "txt2img_steps":
        txt2img_steps_component = component
        for callback in txt2img_steps_callbacks:
            callback(component)


script_callbacks.on_after_component(on_after_component)


def on_ui_settings():
    section = ("freeu", "FreeU")
    shared.opts.add_option(
        "freeu_png_info_auto_enable",
        shared.OptionInfo(
            default=True,
            label="Auto enable when loading the PNG Info of a generation that used FreeU",
            section=section,
        )
    )


script_callbacks.on_ui_settings(on_ui_settings)


unet.patch()
xyz_grid.patch()

scripts/freeu.py
ljleb/sd-webui-freeu