From b14832cf1aa1a9ae68c96820ef1f047e1a9aec50 Mon Sep 17 00:00:00 2001 From: pythongosssss <125205205+pythongosssss@users.noreply.github.com> Date: Sat, 1 Apr 2023 15:10:42 +0100 Subject: [PATCH] Now standard --- tome-patcher/LICENSE | 21 --- tome-patcher/README.md | 5 - tome-patcher/tome_patcher.py | 311 ----------------------------------- 3 files changed, 337 deletions(-) delete mode 100644 tome-patcher/LICENSE delete mode 100644 tome-patcher/README.md delete mode 100644 tome-patcher/tome_patcher.py diff --git a/tome-patcher/LICENSE b/tome-patcher/LICENSE deleted file mode 100644 index ed883e2..0000000 --- a/tome-patcher/LICENSE +++ /dev/null @@ -1,21 +0,0 @@ -MIT License - -Copyright (c) 2023 Daniel Bolya - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. \ No newline at end of file diff --git a/tome-patcher/README.md b/tome-patcher/README.md deleted file mode 100644 index 0d2dcf7..0000000 --- a/tome-patcher/README.md +++ /dev/null @@ -1,5 +0,0 @@ -Adds a TomePatch node that you can use before the KSampler -This will patch the model using tome: -https://github.com/dbolya/tomesd - -It will not currently unpatch the model when it has finished. \ No newline at end of file diff --git a/tome-patcher/tome_patcher.py b/tome-patcher/tome_patcher.py deleted file mode 100644 index 61cde6b..0000000 --- a/tome-patcher/tome_patcher.py +++ /dev/null @@ -1,311 +0,0 @@ -# https://github.com/dbolya/tomesd - -# MIT License - -# Copyright (c) 2023 Daniel Bolya - -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: - -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. - -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -# @article{bolya2023tomesd, -# title={Token Merging for Fast Stable Diffusion}, -# author={Bolya, Daniel and Hoffman, Judy}, -# journal={arXiv}, -# year={2023} -# } - -import torch - - -def isinstance_str(x: object, cls_name: str): - """ - Checks whether x has any class *named* cls_name in its ancestry. - Doesn't require access to the class's implementation. - - Useful for patching! - """ - - for _cls in x.__class__.__mro__: - if _cls.__name__ == cls_name: - return True - - return False -def do_nothing(x: torch.Tensor, mode:str=None): - return x - -import math - -def make_tome_block( - block_class, - ratio: float, - max_downsample: int, - merge_attn: bool, - merge_crossattn: bool, - merge_mlp: bool, - sx: int, sy: int, no_rand: bool): - """ - Make a patched class on the fly so we don't have to import any specific modules. - This patch applies ToMe to the forward function of the block. - """ - - class ToMeBlock(block_class): - # Save for unpatching later - _parent = block_class - - def _forward(self, x: torch.Tensor, context: torch.Tensor = None) -> torch.Tensor: - original_h, original_w = self._tome_info["size"] - original_tokens = original_h * original_w - downsample = int(math.sqrt(original_tokens // x.shape[1])) - - if downsample <= max_downsample: - w = original_w // downsample - h = original_h // downsample - r = int(x.shape[1] * ratio) - m, u = bipartite_soft_matching_random2d(x, w, h, sx, sy, r, no_rand) - else: - m, u = (do_nothing, do_nothing) - - m_a, u_a = (m, u) if merge_attn else (do_nothing, do_nothing) - m_c, u_c = (m, u) if merge_crossattn else (do_nothing, do_nothing) - m_m, u_m = (m, u) if merge_mlp else (do_nothing, do_nothing) - - # This is where the meat of the computation happens - x = u_a(self.attn1(m_a(self.norm1(x)), context=context if self.disable_self_attn else None)) + x - x = u_c(self.attn2(m_c(self.norm2(x)), context=context)) + x - x = u_m(self.ff(m_m(self.norm3(x)))) + x - - return x - - return ToMeBlock - - - -def make_tome_model(model_class): - """ - Make a patched class on the fly so we don't have to import any specific modules. - - This patches the forward function of the model only to get the current image size. - Probably would be better off finding a way to get the size some other way. - """ - - if model_class.__name__ == "ToMeDiffusionModel": - model_class = model_class._parent - - class ToMeDiffusionModel(model_class): - # Save for later - _parent = model_class - - def forward(self, *args, **kwdargs): - self._tome_info["size"] = (args[0].shape[2], args[0].shape[3]) - return super().forward(*args, **kwdargs) - - return ToMeDiffusionModel - -def apply_patch( - model: torch.nn.Module, - ratio: float = 0.5, - max_downsample: int = 1, - sx: int = 2, sy: int = 2, - use_rand: bool = True, - merge_attn: bool = True, - merge_crossattn: bool = False, - merge_mlp: bool = False): - """ - Patches a stable diffusion model with ToMe. - Apply this to the highest level stable diffusion object (i.e., it should have a .model.diffusion_model). - Important Args: - - model: A top level Stable Diffusion module to patch in place. Should have a ".model.diffusion_model" - - ratio: The ratio of tokens to merge. I.e., 0.4 would reduce the total number of tokens by 40%. - The maximum value for this is 1-(1/(sx*sy)). By default, the max is 0.75 (I recommend <= 0.5 though). - Higher values result in more speed-up, but with more visual quality loss. - - Args to tinker with if you want: - - max_downsample [1, 2, 4, or 8]: Apply ToMe to layers with at most this amount of downsampling. - E.g., 1 only applies to layers with no downsampling (4/15) while - 8 applies to all layers (15/15). I recommend a value of 1 or 2. - - sx, sy: The stride for computing dst sets (see paper). A higher stride means you can merge more tokens, - but the default of (2, 2) works well in most cases. Must divide the image size. - - use_rand: Whether or not to allow random perturbations when computing dst sets (see paper). Usually - you'd want to leave this on, but if you're having weird artifacts try turning this off. - - merge_attn: Whether or not to merge tokens for attention (recommended). - - merge_crossattn: Whether or not to merge tokens for cross attention (not recommended). - - merge_mlp: Whether or not to merge tokens for the mlp layers (very not recommended). - """ - - # Make sure the module is not currently patched - remove_patch(model) - - if not hasattr(model, "model") or not hasattr(model.model, "diffusion_model"): - # Provided model not supported - raise RuntimeError("Provided model was not a Stable Diffusion / Latent Diffusion model, as expected.") - - diffusion_model = model.model.diffusion_model - diffusion_model._tome_info = { "size": None, } - diffusion_model.__class__ = make_tome_model(diffusion_model.__class__) - - for _, module in diffusion_model.named_modules(): - # If for some reason this has a different name, create an issue and I'll fix it - if isinstance_str(module, "BasicTransformerBlock"): - module.__class__ = make_tome_block( - module.__class__, ratio, max_downsample, - merge_attn, merge_crossattn, merge_mlp, - sx, sy, not use_rand - ) - module._tome_info = diffusion_model._tome_info - - # Something introduced in SD 2.0 - if not hasattr(module, "disable_self_attn"): - module.disable_self_attn = False - - return model - -def remove_patch(model: torch.nn.Module): - """ Removes a patch from a ToMe Diffusion module if it was already patched. """\ - - for _, module in model.named_modules(): - if module.__class__.__name__ == "ToMeBlock": - module.__class__ = module._parent - elif module.__class__.__name__ == "ToMeDiffusionModel": - module.__class__ = module._parent - - return model - -def bipartite_soft_matching_random2d(metric: torch.Tensor, - w: int, h: int, sx: int, sy: int, r: int, - no_rand: bool = False): - """ - Partitions the tokens into src and dst and merges r tokens from src to dst. - Dst tokens are partitioned by choosing one randomy in each (sx, sy) region. - Args: - - metric [B, N, C]: metric to use for similarity - - w: image width in tokens - - h: image height in tokens - - sx: stride in the x dimension for dst, must divide w - - sy: stride in the y dimension for dst, must divide h - - r: number of tokens to remove (by merging) - - no_rand: if true, disable randomness (use top left corner only) - """ - B, N, _ = metric.shape - - if r <= 0: - return do_nothing, do_nothing - - with torch.no_grad(): - - hsy, wsx = h // sy, w // sx - - # For each sy by sx kernel, randomly assign one token to be dst and the rest src - idx_buffer = torch.zeros(1, hsy, wsx, sy*sx, 1, device=metric.device) - - if no_rand: - rand_idx = torch.zeros(1, hsy, wsx, 1, 1, device=metric.device, dtype=torch.int64) - else: - rand_idx = torch.randint(sy*sx, size=(1, hsy, wsx, 1, 1), device=metric.device) - - idx_buffer.scatter_(dim=3, index=rand_idx, src=-torch.ones_like(rand_idx, dtype=idx_buffer.dtype)) - idx_buffer = idx_buffer.view(1, hsy, wsx, sy, sx, 1).transpose(2, 3).reshape(1, N, 1) - rand_idx = idx_buffer.argsort(dim=1) - - num_dst = int((1 / (sx*sy)) * N) - a_idx = rand_idx[:, num_dst:, :] # src - b_idx = rand_idx[:, :num_dst, :] # dst - - def split(x): - C = x.shape[-1] - src = x.gather(dim=1, index=a_idx.expand(B, N - num_dst, C)) - dst = x.gather(dim=1, index=b_idx.expand(B, num_dst, C)) - return src, dst - - metric = metric / metric.norm(dim=-1, keepdim=True) - a, b = split(metric) - scores = a @ b.transpose(-1, -2) - - # Can't reduce more than the # tokens in src - r = min(a.shape[1], r) - - node_max, node_idx = scores.max(dim=-1) - edge_idx = node_max.argsort(dim=-1, descending=True)[..., None] - - unm_idx = edge_idx[..., r:, :] # Unmerged Tokens - src_idx = edge_idx[..., :r, :] # Merged Tokens - dst_idx = node_idx[..., None].gather(dim=-2, index=src_idx) - - def merge(x: torch.Tensor, mode="mean") -> torch.Tensor: - src, dst = split(x) - n, t1, c = src.shape - - unm = src.gather(dim=-2, index=unm_idx.expand(n, t1 - r, c)) - src = src.gather(dim=-2, index=src_idx.expand(n, r, c)) - dst = dst.scatter_reduce(-2, dst_idx.expand(n, r, c), src, reduce=mode) - - return torch.cat([unm, dst], dim=1) - - def unmerge(x: torch.Tensor) -> torch.Tensor: - unm_len = unm_idx.shape[1] - unm, dst = x[..., :unm_len, :], x[..., unm_len:, :] - _, _, c = unm.shape - - src = dst.gather(dim=-2, index=dst_idx.expand(B, r, c)) - - # Combine back to the original shape - out = torch.zeros(B, N, c, device=x.device, dtype=x.dtype) - out.scatter_(dim=-2, index=b_idx.expand(B, num_dst, c), src=dst) - out.scatter_(dim=-2, index=a_idx.expand(B, a_idx.shape[1], 1).gather(dim=1, index=unm_idx).expand(B, unm_len, c), src=unm) - out.scatter_(dim=-2, index=a_idx.expand(B, a_idx.shape[1], 1).gather(dim=1, index=src_idx).expand(B, r, c), src=src) - - return out - - return merge, unmerge - -class TomePatch: - @classmethod - def INPUT_TYPES(s): - return {"required": { - "model": ("MODEL",), - "ratio": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.1}), - }} - - RETURN_TYPES = ("MODEL",) - FUNCTION = "patch" - - CATEGORY = "utils" - - def patch(self, model, ratio): - model.model = apply_patch(model.model, ratio=ratio) - return (model,) - -# class TomeUnpatch: -# @classmethod -# def INPUT_TYPES(s): -# return {"required": { -# "model": ("MODEL",), -# }} - -# RETURN_TYPES = ("MODEL",) -# FUNCTION = "unpatch" - -# CATEGORY = "utils" - -# def unpatch(self, model, ratio): -# model.model = remove_patch(model.model) -# return (model,) - -NODE_CLASS_MAPPINGS = { - "TomePatch": TomePatch, - # "TomeUnpatch": TomeUnpatch, -}