extensions/sd-webui-pag/scripts/pag.py

import logging
import typing as tg
from collections import OrderedDict
from os import environ
from warnings import warn

import gradio as gr
import torch

import modules.scripts as scripts
from modules import patches, script_callbacks, shared
from modules.processing import StableDiffusionProcessing
from modules.script_callbacks import (
    AfterCFGCallbackParams,
    CFGDenoisedParams,
    CFGDenoiserParams,
)
from modules.sd_samplers_cfg_denoiser import catenate_conds

logger = logging.getLogger(__name__)
logger.setLevel(environ.get("SD_WEBUI_LOG_LEVEL", logging.INFO))

"""
An unofficial implementation of "Self-Rectifying Diffusion Sampling with Perturbed-Attention Guidance" for Automatic1111 WebUI.

@misc{ahn2024selfrectifying,
      title={Self-Rectifying Diffusion Sampling with Perturbed-Attention Guidance},
      author={Donghoon Ahn and Hyoungwon Cho and Jaewon Min and Wooseok Jang and Jungwoo Kim and SeonHwa Kim and Hyun Hee Park and Kyong Hwan Jin and Seungryong Kim},
      year={2024},
      eprint={2403.17377},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

Author: v0xie
GitHub URL: https://github.com/v0xie/sd-webui-incantations

"""


class PAGStateParams:
    def __init__(self) -> None:
        self.pag_scale: float = -1  # PAG guidance scale
        self.guidance_scale: float = -1  # CFG
        self.x_in = None
        self.text_cond: dict | None = None
        self.image_cond: dict | None = None
        self.sigma = None
        self.text_uncond: dict | None = None
        self.make_condition_dict: tg.Callable | None = None  # callable lambda
        self.crossattn_modules: list = []  # callable lambda
        self.to_v_modules: list = []
        self.to_out_modules: list = []
        self.pag_x_out = None
        self.batch_size = -1  # Batch size
        self.denoiser = None  # CFGDenoiser
        self.patched_combine_denoised = None
        self.conds_list = None
        self.uncond_shape_0 = None


class PAGExtensionScript(scripts.Script):
    def __init__(self):
        self.cached_c = [None, None]
        self.handles = []

    # Extension title in menu UI
    def title(self) -> str:
        return "Perturbed Attention Guidance"

    # Decide to show menu in txt2img or img2img
    def show(self, is_img2img):
        return scripts.AlwaysVisible

    # Setup menu ui detail
    def ui(self, is_img2img) -> list:
        with gr.Accordion("Perturbed Attention Guidance", open=False):
            active = gr.Checkbox(
                value=False, default=False, label="Active", elem_id="pag_active"
            )
            with gr.Row():
                pag_scale = gr.Slider(
                    value=3.0,
                    minimum=0,
                    maximum=20.0,
                    step=0.5,
                    label="PAG Scale",
                    elem_id="pag_scale",
                    info="",
                )
        self.infotext_fields = [  # type: ignore
            (active, lambda d: gr.Checkbox.update(value="PAG Active" in d)),
            (pag_scale, "PAG Scale"),
        ]
        self.paste_field_names = [  # type: ignore
            "pag_active",
            "pag_scale",
        ]
        return [active, pag_scale]

    def process_batch(self, p: StableDiffusionProcessing, *args, **kwargs):
        self.pag_process_batch(p, *args, **kwargs)

    def pag_process_batch(
        self, p: StableDiffusionProcessing, active, pag_scale, *args, **kwargs
    ):
        # cleanup previous hooks always
        script_callbacks.remove_current_script_callbacks()
        self.remove_all_hooks()

        active = getattr(p, "pag_active", active)
        if active is False:
            return
        pag_scale = getattr(p, "pag_scale", pag_scale)

        p.extra_generation_params.update(
            {
                "PAG Active": active,
                "PAG Scale": pag_scale,
            }
        )
        self.create_hook(p, active, pag_scale)

    def create_hook(
        self, p: StableDiffusionProcessing, active, pag_scale, *args, **kwargs
    ):
        # Create a list of parameters for each concept
        pag_params = PAGStateParams()
        pag_params.pag_scale = pag_scale
        pag_params.guidance_scale = p.cfg_scale
        pag_params.batch_size = p.batch_size
        pag_params.denoiser = None

        # Get all the qv modules
        cross_attn_modules = self.get_cross_attn_modules()
        if len(cross_attn_modules) == 0:
            logger.error("No cross attention modules found, cannot proceed with PAG")
            return
        pag_params.crossattn_modules = [
            m for m in cross_attn_modules if "CrossAttention" in m.__class__.__name__
        ]

        # Use lambda to call the callback function with the parameters to avoid global variables
        cfg_denoise_lambda = lambda callback_params: self.on_cfg_denoiser_callback(
            callback_params, pag_params
        )
        cfg_denoised_lambda = lambda callback_params: self.on_cfg_denoised_callback(
            callback_params, pag_params
        )
        # after_cfg_lambda = lambda x: self.cfg_after_cfg_callback(x, params)
        unhook_lambda = lambda _: self.unhook_callbacks(pag_params)

        self.ready_hijack_forward(pag_params.crossattn_modules, pag_scale)

        logger.debug("Hooked callbacks")
        script_callbacks.on_cfg_denoiser(cfg_denoise_lambda)
        script_callbacks.on_cfg_denoised(cfg_denoised_lambda)
        # script_callbacks.on_cfg_after_cfg(after_cfg_lambda)
        script_callbacks.on_script_unloaded(unhook_lambda)

    def postprocess_batch(self, p, *args, **kwargs):
        self.pag_postprocess_batch(p, *args, **kwargs)

    def pag_postprocess_batch(self, p, active, *args, **kwargs):
        script_callbacks.remove_current_script_callbacks()

        logger.debug("Removed script callbacks")
        active = getattr(p, "pag_active", active)
        if active is False:
            return

    def remove_all_hooks(self):
        cross_attn_modules = self.get_cross_attn_modules()
        for module in cross_attn_modules:
            to_v = getattr(module, "to_v", None)
            self.remove_field_cross_attn_modules(module, "pag_enable")
            self.remove_field_cross_attn_modules(module, "pag_last_to_v")
            _remove_all_forward_hooks(module, "pag_pre_hook")

            if to_v is not None:
                self.remove_field_cross_attn_modules(to_v, "pag_parent_module")
                _remove_all_forward_hooks(to_v, "to_v_pre_hook")

    def unhook_callbacks(self, pag_params: PAGStateParams):
        if pag_params is None:
            logger.error("PAG params is None")
            return

        if pag_params.denoiser is not None:
            denoiser = pag_params.denoiser
            setattr(denoiser, "combine_denoised_patched", False)
            try:
                patches.undo(__name__, denoiser, "combine_denoised")
            except KeyError:
                logger.exception("KeyError unhooking combine_denoised")
                pass
            except RuntimeError:
                logger.exception("RuntimeError unhooking combine_denoised")
                pass
            pag_params.denoiser = None

    def ready_hijack_forward(self, crossattn_modules, pag_scale):
        """Create hooks in the forward pass of the cross attention modules
        Copies the output of the to_v module to the parent module
        Then applies the PAG perturbation to the output of the cross attention module (multiplication by identity)
        """

        # add field for last_to_v
        for module in crossattn_modules:
            to_v = getattr(module, "to_v", None)
            self.add_field_cross_attn_modules(module, "pag_enable", False)
            self.add_field_cross_attn_modules(module, "pag_last_to_v", None)
            self.add_field_cross_attn_modules(to_v, "pag_parent_module", [module])
            # self.add_field_cross_attn_modules(to_out, 'pag_parent_module', [module])

        def to_v_pre_hook(module, input, kwargs, output):
            """Copy the output of the to_v module to the parent module"""
            parent_module = getattr(module, "pag_parent_module", None)
            # copy the output of the to_v module to the parent module

            if parent_module is not None:
                setattr(parent_module[0], "pag_last_to_v", output.detach().clone())

        def pag_pre_hook(module, input, kwargs, output):
            if (
                hasattr(module, "pag_enable")
                and getattr(module, "pag_enable", False) is False
            ):
                return
            if not hasattr(module, "pag_last_to_v"):
                # oops we forgot to unhook
                return

            batch_size, seq_len, inner_dim = output.shape
            identity = torch.eye(seq_len).expand(batch_size, -1, -1).to(shared.device)

            # get the last to_v output and save it
            last_to_v = getattr(module, "pag_last_to_v", None)
            if last_to_v is not None:
                new_output = torch.einsum("bij,bjk->bik", identity, last_to_v)
                return new_output
            else:
                # this is bad
                return output

        # Create hooks
        for module in crossattn_modules:
            module.register_forward_hook(pag_pre_hook, with_kwargs=True)

            to_v = getattr(module, "to_v", None)
            if to_v is not None:
                to_v.register_forward_hook(to_v_pre_hook, with_kwargs=True)

    def get_middle_block_modules(self):
        """Get all attention modules from the middle block
        Refere to page 22 of the PAG paper, Appendix A.2

        """
        try:
            m = shared.sd_model
            nlm = m.network_layer_mapping
            middle_block_modules = [
                m
                for m in nlm.values()
                if "middle_block_1_transformer_blocks_0_attn1" in m.network_layer_name
                and "CrossAttention" in m.__class__.__name__
            ]
            return middle_block_modules
        except AttributeError:
            logger.exception(
                "AttributeError in get_middle_block_modules", stack_info=True
            )
            return []
        except Exception:
            logger.exception("Exception in get_middle_block_modules", stack_info=True)
            return []

    def get_cross_attn_modules(self):
        """Get all cross attention modules"""
        return self.get_middle_block_modules()

    def add_field_cross_attn_modules(self, module, field, value):
        """Add a field to a module if it doesn't exist"""
        if not hasattr(module, field):
            setattr(module, field, value)

    def remove_field_cross_attn_modules(self, module, field):
        """Remove a field from a module if it exists"""
        if hasattr(module, field):
            delattr(module, field)

    def on_cfg_denoiser_callback(
        self, params: CFGDenoiserParams, pag_params: PAGStateParams
    ):
        # always unhook
        self.unhook_callbacks(pag_params)

        # patch combine_denoised
        if pag_params.denoiser is None:
            pag_params.denoiser = params.denoiser
        if getattr(params.denoiser, "combine_denoised_patched", False) is False:
            try:
                setattr(
                    params.denoiser,
                    "combine_denoised_original",
                    getattr(params.denoiser, "combine_denoised"),
                )
                # create patch that references the original function
                pass_conds_func = (
                    lambda *args, **kwargs: combine_denoised_pass_conds_list(
                        *args,
                        **kwargs,
                        original_func=getattr(
                            params.denoiser, "combine_denoised_original"
                        ),
                        pag_params=pag_params
                    )
                )
                pag_params.patched_combine_denoised = patches.patch(
                    __name__, params.denoiser, "combine_denoised", pass_conds_func
                )
                setattr(params.denoiser, "combine_denoised_patched", True)
                setattr(
                    params.denoiser,
                    "combine_denoised_original",
                    patches.original(__name__, params.denoiser, "combine_denoised"),
                )
            except KeyError:
                logger.exception("KeyError patching combine_denoised")
                pass
            except RuntimeError:
                logger.exception("RuntimeError patching combine_denoised")
                pass

        if isinstance(params.text_cond, dict):
            text_cond = params.text_cond["crossattn"]  # SD XL
            pag_params.text_cond = {}
            pag_params.text_uncond = {}
            for key, value in params.text_cond.items():
                pag_params.text_cond[key] = value.clone().detach()
                pag_params.text_uncond[key] = value.clone().detach()
        else:
            text_cond = params.text_cond  # SD 1.5
            pag_params.text_cond = text_cond.clone().detach()
            pag_params.text_uncond = text_cond.clone().detach()

        pag_params.x_in = params.x.clone().detach()
        pag_params.sigma = params.sigma.clone().detach()
        pag_params.image_cond = params.image_cond.clone().detach()
        pag_params.denoiser = params.denoiser
        pag_params.make_condition_dict = get_make_condition_dict_fn(params.text_uncond)

    def on_cfg_denoised_callback(
        self, params: CFGDenoisedParams, pag_params: PAGStateParams
    ):
        """Callback function for the CFGDenoisedParams
        Refer to pg.22 A.2 of the PAG paper for how CFG and PAG combine

        """
        # passed from on_cfg_denoiser_callback
        x_in = pag_params.x_in
        tensor = pag_params.text_cond
        uncond = pag_params.text_uncond
        image_cond_in = pag_params.image_cond
        sigma_in = pag_params.sigma

        # concatenate the conditions
        # "modules/sd_samplers_cfg_denoiser.py:237"
        cond_in = catenate_conds([tensor, uncond])
        make_condition_dict = get_make_condition_dict_fn(uncond)
        conds = make_condition_dict(cond_in, image_cond_in)

        # set pag_enable to True for the hooked cross attention modules
        for module in pag_params.crossattn_modules:
            setattr(module, "pag_enable", True)

        # get the PAG guidance (is there a way to optimize this so we don't have to calculate it twice?)
        pag_x_out = params.inner_model(x_in, sigma_in, cond=conds)

        # update pag_x_out
        pag_params.pag_x_out = pag_x_out

        # set pag_enable to False
        for module in pag_params.crossattn_modules:
            setattr(module, "pag_enable", False)

    def cfg_after_cfg_callback(
        self, params: AfterCFGCallbackParams, pag_params: PAGStateParams
    ):
        # self.unhook_callbacks(pag_params)
        pass


def combine_denoised_pass_conds_list(*args, **kwargs):
    """Hijacked function for combine_denoised in CFGDenoiser"""
    original_func = kwargs.get("original_func", None)
    new_params = kwargs.get("pag_params", None)

    if new_params is None:
        logger.error("new_params is None")
        return original_func(*args)

    def new_combine_denoised(x_out, conds_list, uncond, cond_scale):
        denoised_uncond = x_out[-uncond.shape[0] :]
        denoised = torch.clone(denoised_uncond)

        for i, conds in enumerate(conds_list):
            for cond_index, weight in conds:
                denoised[i] += (x_out[cond_index] - denoised_uncond[i]) * (
                    weight * cond_scale
                )
                try:
                    denoised[i] += (x_out[cond_index] - new_params.pag_x_out[i]) * (
                        weight * new_params.pag_scale
                    )
                except TypeError:
                    logger.exception("TypeError in combine_denoised_pass_conds_list")
                except IndexError:
                    logger.exception("IndexError in combine_denoised_pass_conds_list")
                # logger.debug(f"added PAG guidance to denoised - pag_scale:{global_scale}")
        return denoised

    return new_combine_denoised(*args)


# from modules/sd_samplers_cfg_denoiser.py:187-195
def get_make_condition_dict_fn(text_uncond):
    if shared.sd_model.model.conditioning_key == "crossattn-adm":
        make_condition_dict = lambda c_crossattn, c_adm: {
            "c_crossattn": [c_crossattn],
            "c_adm": c_adm,
        }
    else:
        if isinstance(text_uncond, dict):
            make_condition_dict = lambda c_crossattn, c_concat: {
                **c_crossattn,
                "c_concat": [c_concat],
            }
        else:
            make_condition_dict = lambda c_crossattn, c_concat: {
                "c_crossattn": [c_crossattn],
                "c_concat": [c_concat],
            }
    return make_condition_dict


# thanks torch; removing hooks DOESN'T WORK
# thank you to @ProGamerGov for this https://github.com/pytorch/pytorch/issues/70455
def _remove_all_forward_hooks(
    module: torch.nn.Module, hook_fn_name: str | None = None
) -> None:
    """
    This function removes all forward hooks in the specified module, without requiring
    any hook handles. This lets us clean up & remove any hooks that weren't property
    deleted.

    Warning: Various PyTorch modules and systems make use of hooks, and thus extreme
    caution should be exercised when removing all hooks. Users are recommended to give
    their hook function a unique name that can be used to safely identify and remove
    the target forward hooks.

    Args:

        module (nn.Module): The module instance to remove forward hooks from.
        hook_fn_name (str, optional): Optionally only remove specific forward hooks
            based on their function's __name__ attribute.
            Default: None
    """

    if hook_fn_name is None:
        warn("Removing all active hooks can break some PyTorch modules & systems.")

    def _remove_hooks(m: torch.nn.Module, name: str | None = None) -> None:
        if hasattr(module, "_forward_hooks"):
            if m._forward_hooks != OrderedDict():
                if name is not None:
                    dict_items = list(m._forward_hooks.items())
                    m._forward_hooks = OrderedDict(
                        [(i, fn) for i, fn in dict_items if fn.__name__ != name]
                    )
                else:
                    m._forward_hooks = OrderedDict()

    def _remove_child_hooks(
        target_module: torch.nn.Module, hook_name: str | None = None
    ) -> None:
        for name, child in target_module._modules.items():
            if child is not None:
                _remove_hooks(child, hook_name)
                _remove_child_hooks(child, hook_name)

    # Remove hooks from target submodules
    _remove_child_hooks(module, hook_fn_name)

    # Remove hooks from the target module
    _remove_hooks(module, hook_fn_name)