Merge branch 'AUTOMATIC1111:master' into master

8 files changed, 670 insertions, 9 deletions

diff --git a/extensions-builtin/Lora/lyco_helpers.py b/extensions-builtin/Lora/lyco_helpers.py
index 279b34bc..1679a0ce 100644
--- a/extensions-builtin/Lora/lyco_helpers.py
+++ b/extensions-builtin/Lora/lyco_helpers.py
@@ -19,3 +19,50 @@ def rebuild_cp_decomposition(up, down, mid):
     up = up.reshape(up.size(0), -1)

     down = down.reshape(down.size(0), -1)

     return torch.einsum('n m k l, i n, m j -> i j k l', mid, up, down)

+

+

+# copied from https://github.com/KohakuBlueleaf/LyCORIS/blob/dev/lycoris/modules/lokr.py

+def factorization(dimension: int, factor:int=-1) -> tuple[int, int]:

+    '''

+    return a tuple of two value of input dimension decomposed by the number closest to factor

+    second value is higher or equal than first value.

+

+    In LoRA with Kroneckor Product, first value is a value for weight scale.

+    secon value is a value for weight.

+

+    Becuase of non-commutative property, A⊗B ≠ B⊗A. Meaning of two matrices is slightly different.

+

+    examples)

+    factor

+        -1               2                4               8               16               ...

+    127 -> 1, 127   127 -> 1, 127    127 -> 1, 127   127 -> 1, 127   127 -> 1, 127

+    128 -> 8, 16    128 -> 2, 64     128 -> 4, 32    128 -> 8, 16    128 -> 8, 16

+    250 -> 10, 25   250 -> 2, 125    250 -> 2, 125   250 -> 5, 50    250 -> 10, 25

+    360 -> 8, 45    360 -> 2, 180    360 -> 4, 90    360 -> 8, 45    360 -> 12, 30

+    512 -> 16, 32   512 -> 2, 256    512 -> 4, 128   512 -> 8, 64    512 -> 16, 32

+    1024 -> 32, 32  1024 -> 2, 512   1024 -> 4, 256  1024 -> 8, 128  1024 -> 16, 64

+    '''

+

+    if factor > 0 and (dimension % factor) == 0:

+        m = factor

+        n = dimension // factor

+        if m > n:

+            n, m = m, n

+        return m, n

+    if factor < 0:

+        factor = dimension

+    m, n = 1, dimension

+    length = m + n

+    while m<n:

+        new_m = m + 1

+        while dimension%new_m != 0:

+            new_m += 1

+        new_n = dimension // new_m

+        if new_m + new_n > length or new_m>factor:

+            break

+        else:

+            m, n = new_m, new_n

+    if m > n:

+        n, m = m, n

+    return m, n

+

diff --git a/extensions-builtin/Lora/network_oft.py b/extensions-builtin/Lora/network_oft.py
new file mode 100644
index 00000000..fa647020
--- /dev/null
+++ b/extensions-builtin/Lora/network_oft.py
@@ -0,0 +1,82 @@
+import torch
+import network
+from lyco_helpers import factorization
+from einops import rearrange
+
+
+class ModuleTypeOFT(network.ModuleType):
+    def create_module(self, net: network.Network, weights: network.NetworkWeights):
+        if all(x in weights.w for x in ["oft_blocks"]) or all(x in weights.w for x in ["oft_diag"]):
+            return NetworkModuleOFT(net, weights)
+
+        return None
+
+# Supports both kohya-ss' implementation of COFT  https://github.com/kohya-ss/sd-scripts/blob/main/networks/oft.py
+# and KohakuBlueleaf's implementation of OFT/COFT https://github.com/KohakuBlueleaf/LyCORIS/blob/dev/lycoris/modules/diag_oft.py
+class NetworkModuleOFT(network.NetworkModule):
+    def __init__(self,  net: network.Network, weights: network.NetworkWeights):
+
+        super().__init__(net, weights)
+
+        self.lin_module = None
+        self.org_module: list[torch.Module] = [self.sd_module]
+
+        self.scale = 1.0
+
+        # kohya-ss
+        if "oft_blocks" in weights.w.keys():
+            self.is_kohya = True
+            self.oft_blocks = weights.w["oft_blocks"] # (num_blocks, block_size, block_size)
+            self.alpha = weights.w["alpha"] # alpha is constraint
+            self.dim = self.oft_blocks.shape[0] # lora dim
+        # LyCORIS
+        elif "oft_diag" in weights.w.keys():
+            self.is_kohya = False
+            self.oft_blocks = weights.w["oft_diag"]
+            # self.alpha is unused
+            self.dim = self.oft_blocks.shape[1] # (num_blocks, block_size, block_size)
+
+        is_linear = type(self.sd_module) in [torch.nn.Linear, torch.nn.modules.linear.NonDynamicallyQuantizableLinear]
+        is_conv = type(self.sd_module) in [torch.nn.Conv2d]
+        is_other_linear = type(self.sd_module) in [torch.nn.MultiheadAttention] # unsupported
+
+        if is_linear:
+            self.out_dim = self.sd_module.out_features
+        elif is_conv:
+            self.out_dim = self.sd_module.out_channels
+        elif is_other_linear:
+            self.out_dim = self.sd_module.embed_dim
+
+        if self.is_kohya:
+            self.constraint = self.alpha * self.out_dim
+            self.num_blocks = self.dim
+            self.block_size = self.out_dim // self.dim
+        else:
+            self.constraint = None
+            self.block_size, self.num_blocks = factorization(self.out_dim, self.dim)
+
+    def calc_updown(self, orig_weight):
+        oft_blocks = self.oft_blocks.to(orig_weight.device, dtype=orig_weight.dtype)
+        eye = torch.eye(self.block_size, device=self.oft_blocks.device)
+
+        if self.is_kohya:
+            block_Q = oft_blocks - oft_blocks.transpose(1, 2) # ensure skew-symmetric orthogonal matrix
+            norm_Q = torch.norm(block_Q.flatten())
+            new_norm_Q = torch.clamp(norm_Q, max=self.constraint)
+            block_Q = block_Q * ((new_norm_Q + 1e-8) / (norm_Q + 1e-8))
+            oft_blocks = torch.matmul(eye + block_Q, (eye - block_Q).float().inverse())
+
+        R = oft_blocks.to(orig_weight.device, dtype=orig_weight.dtype)
+
+        # This errors out for MultiheadAttention, might need to be handled up-stream
+        merged_weight = rearrange(orig_weight, '(k n) ... -> k n ...', k=self.num_blocks, n=self.block_size)
+        merged_weight = torch.einsum(
+            'k n m, k n ... -> k m ...',
+            R,
+            merged_weight
+        )
+        merged_weight = rearrange(merged_weight, 'k m ... -> (k m) ...')
+
+        updown = merged_weight.to(orig_weight.device, dtype=orig_weight.dtype) - orig_weight
+        output_shape = orig_weight.shape
+        return self.finalize_updown(updown, orig_weight, output_shape)
diff --git a/extensions-builtin/Lora/networks.py b/extensions-builtin/Lora/networks.py
index 60d8dec4..629bf853 100644
--- a/extensions-builtin/Lora/networks.py
+++ b/extensions-builtin/Lora/networks.py
@@ -11,6 +11,7 @@ import network_ia3
 import network_lokr

 import network_full

 import network_norm

+import network_oft

 

 import torch

 from typing import Union

@@ -28,6 +29,7 @@ module_types = [
     network_full.ModuleTypeFull(),

     network_norm.ModuleTypeNorm(),

     network_glora.ModuleTypeGLora(),

+    network_oft.ModuleTypeOFT(),

 ]

 

 

@@ -157,7 +159,8 @@ def load_network(name, network_on_disk):
     bundle_embeddings = {}

 

     for key_network, weight in sd.items():

-        key_network_without_network_parts, network_part = key_network.split(".", 1)

+        key_network_without_network_parts, _, network_part = key_network.partition(".")

+

         if key_network_without_network_parts == "bundle_emb":

             emb_name, vec_name = network_part.split(".", 1)

             emb_dict = bundle_embeddings.get(emb_name, {})

@@ -189,6 +192,17 @@ def load_network(name, network_on_disk):
                 key = key_network_without_network_parts.replace("lora_te1_text_model", "transformer_text_model")

                 sd_module = shared.sd_model.network_layer_mapping.get(key, None)

 

+        # kohya_ss OFT module

+        elif sd_module is None and "oft_unet" in key_network_without_network_parts:

+            key = key_network_without_network_parts.replace("oft_unet", "diffusion_model")

+            sd_module = shared.sd_model.network_layer_mapping.get(key, None)

+

+        # KohakuBlueLeaf OFT module

+        if sd_module is None and "oft_diag" in key:

+            key = key_network_without_network_parts.replace("lora_unet", "diffusion_model")

+            key = key_network_without_network_parts.replace("lora_te1_text_model", "0_transformer_text_model")

+            sd_module = shared.sd_model.network_layer_mapping.get(key, None)

+

         if sd_module is None:

             keys_failed_to_match[key_network] = key

             continue

diff --git a/extensions-builtin/Lora/ui_extra_networks_lora.py b/extensions-builtin/Lora/ui_extra_networks_lora.py
index 55409a78..df02c663 100644
--- a/extensions-builtin/Lora/ui_extra_networks_lora.py
+++ b/extensions-builtin/Lora/ui_extra_networks_lora.py
@@ -17,6 +17,8 @@ class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage):
 

     def create_item(self, name, index=None, enable_filter=True):

         lora_on_disk = networks.available_networks.get(name)

+        if lora_on_disk is None:

+            return

 

         path, ext = os.path.splitext(lora_on_disk.filename)

 

@@ -66,9 +68,10 @@ class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage):
         return item

 

     def list_items(self):

-        for index, name in enumerate(networks.available_networks):

+        # instantiate a list to protect against concurrent modification

+        names = list(networks.available_networks)

+        for index, name in enumerate(names):

             item = self.create_item(name, index)

-

             if item is not None:

                 yield item

 

diff --git a/extensions-builtin/extra-options-section/scripts/extra_options_section.py b/extensions-builtin/extra-options-section/scripts/extra_options_section.py
index 983f87ff..ac2c3de4 100644
--- a/extensions-builtin/extra-options-section/scripts/extra_options_section.py
+++ b/extensions-builtin/extra-options-section/scripts/extra_options_section.py
@@ -23,11 +23,12 @@ class ExtraOptionsSection(scripts.Script):
         self.setting_names = []

         self.infotext_fields = []

         extra_options = shared.opts.extra_options_img2img if is_img2img else shared.opts.extra_options_txt2img

+        elem_id_tabname = "extra_options_" + ("img2img" if is_img2img else "txt2img")

 

         mapping = {k: v for v, k in generation_parameters_copypaste.infotext_to_setting_name_mapping}

 

         with gr.Blocks() as interface:

-            with gr.Accordion("Options", open=False) if shared.opts.extra_options_accordion and extra_options else gr.Group():

+            with gr.Accordion("Options", open=False, elem_id=elem_id_tabname) if shared.opts.extra_options_accordion and extra_options else gr.Group(elem_id=elem_id_tabname):

 

                 row_count = math.ceil(len(extra_options) / shared.opts.extra_options_cols)

 

@@ -64,11 +65,14 @@ class ExtraOptionsSection(scripts.Script):
                 p.override_settings[name] = value

 

 

-shared.options_templates.update(shared.options_section(('ui', "User interface"), {

-    "extra_options_txt2img": shared.OptionInfo([], "Options in main UI - txt2img", ui_components.DropdownMulti, lambda: {"choices": list(shared.opts.data_labels.keys())}).js("info", "settingsHintsShowQuicksettings").info("setting entries that also appear in txt2img interfaces").needs_reload_ui(),

-    "extra_options_img2img": shared.OptionInfo([], "Options in main UI - img2img", ui_components.DropdownMulti, lambda: {"choices": list(shared.opts.data_labels.keys())}).js("info", "settingsHintsShowQuicksettings").info("setting entries that also appear in img2img interfaces").needs_reload_ui(),

-    "extra_options_cols": shared.OptionInfo(1, "Options in main UI - number of columns", gr.Number, {"precision": 0}).needs_reload_ui(),

-    "extra_options_accordion": shared.OptionInfo(False, "Options in main UI - place into an accordion").needs_reload_ui()

+shared.options_templates.update(shared.options_section(('settings_in_ui', "Settings in UI", "ui"), {

+    "settings_in_ui": shared.OptionHTML("""

+This page allows you to add some settings to the main interface of txt2img and img2img tabs.

+"""),

+    "extra_options_txt2img": shared.OptionInfo([], "Settings for txt2img", ui_components.DropdownMulti, lambda: {"choices": list(shared.opts.data_labels.keys())}).js("info", "settingsHintsShowQuicksettings").info("setting entries that also appear in txt2img interfaces").needs_reload_ui(),

+    "extra_options_img2img": shared.OptionInfo([], "Settings for img2img", ui_components.DropdownMulti, lambda: {"choices": list(shared.opts.data_labels.keys())}).js("info", "settingsHintsShowQuicksettings").info("setting entries that also appear in img2img interfaces").needs_reload_ui(),

+    "extra_options_cols": shared.OptionInfo(1, "Number of columns for added settings", gr.Slider, {"step": 1, "minimum": 1, "maximum": 20}).info("displayed amount will depend on the actual browser window width").needs_reload_ui(),

+    "extra_options_accordion": shared.OptionInfo(False, "Place added settings into an accordion").needs_reload_ui()

 }))

 

 

diff --git a/extensions-builtin/hypertile/hypertile.py b/extensions-builtin/hypertile/hypertile.py
new file mode 100644
index 00000000..0f40e2d3
--- /dev/null
+++ b/extensions-builtin/hypertile/hypertile.py
@@ -0,0 +1,351 @@
+"""
+Hypertile module for splitting attention layers in SD-1.5 U-Net and SD-1.5 VAE
+Warn: The patch works well only if the input image has a width and height that are multiples of 128
+Original author: @tfernd Github: https://github.com/tfernd/HyperTile
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable
+
+from functools import wraps, cache
+
+import math
+import torch.nn as nn
+import random
+
+from einops import rearrange
+
+
+@dataclass
+class HypertileParams:
+    depth = 0
+    layer_name = ""
+    tile_size: int = 0
+    swap_size: int = 0
+    aspect_ratio: float = 1.0
+    forward = None
+    enabled = False
+
+
+
+# TODO add SD-XL layers
+DEPTH_LAYERS = {
+    0: [
+        # SD 1.5 U-Net (diffusers)
+        "down_blocks.0.attentions.0.transformer_blocks.0.attn1",
+        "down_blocks.0.attentions.1.transformer_blocks.0.attn1",
+        "up_blocks.3.attentions.0.transformer_blocks.0.attn1",
+        "up_blocks.3.attentions.1.transformer_blocks.0.attn1",
+        "up_blocks.3.attentions.2.transformer_blocks.0.attn1",
+        # SD 1.5 U-Net (ldm)
+        "input_blocks.1.1.transformer_blocks.0.attn1",
+        "input_blocks.2.1.transformer_blocks.0.attn1",
+        "output_blocks.9.1.transformer_blocks.0.attn1",
+        "output_blocks.10.1.transformer_blocks.0.attn1",
+        "output_blocks.11.1.transformer_blocks.0.attn1",
+        # SD 1.5 VAE
+        "decoder.mid_block.attentions.0",
+        "decoder.mid.attn_1",
+    ],
+    1: [
+        # SD 1.5 U-Net (diffusers)
+        "down_blocks.1.attentions.0.transformer_blocks.0.attn1",
+        "down_blocks.1.attentions.1.transformer_blocks.0.attn1",
+        "up_blocks.2.attentions.0.transformer_blocks.0.attn1",
+        "up_blocks.2.attentions.1.transformer_blocks.0.attn1",
+        "up_blocks.2.attentions.2.transformer_blocks.0.attn1",
+        # SD 1.5 U-Net (ldm)
+        "input_blocks.4.1.transformer_blocks.0.attn1",
+        "input_blocks.5.1.transformer_blocks.0.attn1",
+        "output_blocks.6.1.transformer_blocks.0.attn1",
+        "output_blocks.7.1.transformer_blocks.0.attn1",
+        "output_blocks.8.1.transformer_blocks.0.attn1",
+    ],
+    2: [
+        # SD 1.5 U-Net (diffusers)
+        "down_blocks.2.attentions.0.transformer_blocks.0.attn1",
+        "down_blocks.2.attentions.1.transformer_blocks.0.attn1",
+        "up_blocks.1.attentions.0.transformer_blocks.0.attn1",
+        "up_blocks.1.attentions.1.transformer_blocks.0.attn1",
+        "up_blocks.1.attentions.2.transformer_blocks.0.attn1",
+        # SD 1.5 U-Net (ldm)
+        "input_blocks.7.1.transformer_blocks.0.attn1",
+        "input_blocks.8.1.transformer_blocks.0.attn1",
+        "output_blocks.3.1.transformer_blocks.0.attn1",
+        "output_blocks.4.1.transformer_blocks.0.attn1",
+        "output_blocks.5.1.transformer_blocks.0.attn1",
+    ],
+    3: [
+        # SD 1.5 U-Net (diffusers)
+        "mid_block.attentions.0.transformer_blocks.0.attn1",
+        # SD 1.5 U-Net (ldm)
+        "middle_block.1.transformer_blocks.0.attn1",
+    ],
+}
+# XL layers, thanks for GitHub@gel-crabs for the help
+DEPTH_LAYERS_XL = {
+    0: [
+        # SD 1.5 U-Net (diffusers)
+        "down_blocks.0.attentions.0.transformer_blocks.0.attn1",
+        "down_blocks.0.attentions.1.transformer_blocks.0.attn1",
+        "up_blocks.3.attentions.0.transformer_blocks.0.attn1",
+        "up_blocks.3.attentions.1.transformer_blocks.0.attn1",
+        "up_blocks.3.attentions.2.transformer_blocks.0.attn1",
+        # SD 1.5 U-Net (ldm)
+        "input_blocks.4.1.transformer_blocks.0.attn1",
+        "input_blocks.5.1.transformer_blocks.0.attn1",
+        "output_blocks.3.1.transformer_blocks.0.attn1",
+        "output_blocks.4.1.transformer_blocks.0.attn1",
+        "output_blocks.5.1.transformer_blocks.0.attn1",
+        # SD 1.5 VAE
+        "decoder.mid_block.attentions.0",
+        "decoder.mid.attn_1",
+    ],
+    1: [
+        # SD 1.5 U-Net (diffusers)
+        #"down_blocks.1.attentions.0.transformer_blocks.0.attn1",
+        #"down_blocks.1.attentions.1.transformer_blocks.0.attn1",
+        #"up_blocks.2.attentions.0.transformer_blocks.0.attn1",
+        #"up_blocks.2.attentions.1.transformer_blocks.0.attn1",
+        #"up_blocks.2.attentions.2.transformer_blocks.0.attn1",
+        # SD 1.5 U-Net (ldm)
+        "input_blocks.4.1.transformer_blocks.1.attn1",
+        "input_blocks.5.1.transformer_blocks.1.attn1",
+        "output_blocks.3.1.transformer_blocks.1.attn1",
+        "output_blocks.4.1.transformer_blocks.1.attn1",
+        "output_blocks.5.1.transformer_blocks.1.attn1",
+        "input_blocks.7.1.transformer_blocks.0.attn1",
+        "input_blocks.8.1.transformer_blocks.0.attn1",
+        "output_blocks.0.1.transformer_blocks.0.attn1",
+        "output_blocks.1.1.transformer_blocks.0.attn1",
+        "output_blocks.2.1.transformer_blocks.0.attn1",
+        "input_blocks.7.1.transformer_blocks.1.attn1",
+        "input_blocks.8.1.transformer_blocks.1.attn1",
+        "output_blocks.0.1.transformer_blocks.1.attn1",
+        "output_blocks.1.1.transformer_blocks.1.attn1",
+        "output_blocks.2.1.transformer_blocks.1.attn1",
+        "input_blocks.7.1.transformer_blocks.2.attn1",
+        "input_blocks.8.1.transformer_blocks.2.attn1",
+        "output_blocks.0.1.transformer_blocks.2.attn1",
+        "output_blocks.1.1.transformer_blocks.2.attn1",
+        "output_blocks.2.1.transformer_blocks.2.attn1",
+        "input_blocks.7.1.transformer_blocks.3.attn1",
+        "input_blocks.8.1.transformer_blocks.3.attn1",
+        "output_blocks.0.1.transformer_blocks.3.attn1",
+        "output_blocks.1.1.transformer_blocks.3.attn1",
+        "output_blocks.2.1.transformer_blocks.3.attn1",
+        "input_blocks.7.1.transformer_blocks.4.attn1",
+        "input_blocks.8.1.transformer_blocks.4.attn1",
+        "output_blocks.0.1.transformer_blocks.4.attn1",
+        "output_blocks.1.1.transformer_blocks.4.attn1",
+        "output_blocks.2.1.transformer_blocks.4.attn1",
+        "input_blocks.7.1.transformer_blocks.5.attn1",
+        "input_blocks.8.1.transformer_blocks.5.attn1",
+        "output_blocks.0.1.transformer_blocks.5.attn1",
+        "output_blocks.1.1.transformer_blocks.5.attn1",
+        "output_blocks.2.1.transformer_blocks.5.attn1",
+        "input_blocks.7.1.transformer_blocks.6.attn1",
+        "input_blocks.8.1.transformer_blocks.6.attn1",
+        "output_blocks.0.1.transformer_blocks.6.attn1",
+        "output_blocks.1.1.transformer_blocks.6.attn1",
+        "output_blocks.2.1.transformer_blocks.6.attn1",
+        "input_blocks.7.1.transformer_blocks.7.attn1",
+        "input_blocks.8.1.transformer_blocks.7.attn1",
+        "output_blocks.0.1.transformer_blocks.7.attn1",
+        "output_blocks.1.1.transformer_blocks.7.attn1",
+        "output_blocks.2.1.transformer_blocks.7.attn1",
+        "input_blocks.7.1.transformer_blocks.8.attn1",
+        "input_blocks.8.1.transformer_blocks.8.attn1",
+        "output_blocks.0.1.transformer_blocks.8.attn1",
+        "output_blocks.1.1.transformer_blocks.8.attn1",
+        "output_blocks.2.1.transformer_blocks.8.attn1",
+        "input_blocks.7.1.transformer_blocks.9.attn1",
+        "input_blocks.8.1.transformer_blocks.9.attn1",
+        "output_blocks.0.1.transformer_blocks.9.attn1",
+        "output_blocks.1.1.transformer_blocks.9.attn1",
+        "output_blocks.2.1.transformer_blocks.9.attn1",
+    ],
+    2: [
+        # SD 1.5 U-Net (diffusers)
+        "mid_block.attentions.0.transformer_blocks.0.attn1",
+        # SD 1.5 U-Net (ldm)
+        "middle_block.1.transformer_blocks.0.attn1",
+        "middle_block.1.transformer_blocks.1.attn1",
+        "middle_block.1.transformer_blocks.2.attn1",
+        "middle_block.1.transformer_blocks.3.attn1",
+        "middle_block.1.transformer_blocks.4.attn1",
+        "middle_block.1.transformer_blocks.5.attn1",
+        "middle_block.1.transformer_blocks.6.attn1",
+        "middle_block.1.transformer_blocks.7.attn1",
+        "middle_block.1.transformer_blocks.8.attn1",
+        "middle_block.1.transformer_blocks.9.attn1",
+    ],
+    3 : [] # TODO - separate layers for SD-XL
+}
+
+
+RNG_INSTANCE = random.Random()
+
+@cache
+def get_divisors(value: int, min_value: int, /, max_options: int = 1) -> list[int]:
+    """
+    Returns divisors of value that
+        x * min_value <= value
+    in big -> small order, amount of divisors is limited by max_options
+    """
+    max_options = max(1, max_options) # at least 1 option should be returned
+    min_value = min(min_value, value)
+    divisors = [i for i in range(min_value, value + 1) if value % i == 0] # divisors in small -> big order
+    ns = [value // i for i in divisors[:max_options]]  # has at least 1 element # big -> small order
+    return ns
+
+
+def random_divisor(value: int, min_value: int, /, max_options: int = 1) -> int:
+    """
+    Returns a random divisor of value that
+        x * min_value <= value
+    if max_options is 1, the behavior is deterministic
+    """
+    ns = get_divisors(value, min_value, max_options=max_options) # get cached divisors
+    idx = RNG_INSTANCE.randint(0, len(ns) - 1)
+
+    return ns[idx]
+
+
+def set_hypertile_seed(seed: int) -> None:
+    RNG_INSTANCE.seed(seed)
+
+
+@cache
+def largest_tile_size_available(width: int, height: int) -> int:
+    """
+    Calculates the largest tile size available for a given width and height
+    Tile size is always a power of 2
+    """
+    gcd = math.gcd(width, height)
+    largest_tile_size_available = 1
+    while gcd % (largest_tile_size_available * 2) == 0:
+        largest_tile_size_available *= 2
+    return largest_tile_size_available
+
+
+def iterative_closest_divisors(hw:int, aspect_ratio:float) -> tuple[int, int]:
+    """
+    Finds h and w such that h*w = hw and h/w = aspect_ratio
+    We check all possible divisors of hw and return the closest to the aspect ratio
+    """
+    divisors = [i for i in range(2, hw + 1) if hw % i == 0] # all divisors of hw
+    pairs = [(i, hw // i) for i in divisors] # all pairs of divisors of hw
+    ratios = [w/h for h, w in pairs] # all ratios of pairs of divisors of hw
+    closest_ratio = min(ratios, key=lambda x: abs(x - aspect_ratio)) # closest ratio to aspect_ratio
+    closest_pair = pairs[ratios.index(closest_ratio)] # closest pair of divisors to aspect_ratio
+    return closest_pair
+
+
+@cache
+def find_hw_candidates(hw:int, aspect_ratio:float) -> tuple[int, int]:
+    """
+    Finds h and w such that h*w = hw and h/w = aspect_ratio
+    """
+    h, w = round(math.sqrt(hw * aspect_ratio)), round(math.sqrt(hw / aspect_ratio))
+    # find h and w such that h*w = hw and h/w = aspect_ratio
+    if h * w != hw:
+        w_candidate = hw / h
+        # check if w is an integer
+        if not w_candidate.is_integer():
+            h_candidate = hw / w
+            # check if h is an integer
+            if not h_candidate.is_integer():
+                return iterative_closest_divisors(hw, aspect_ratio)
+            else:
+                h = int(h_candidate)
+        else:
+            w = int(w_candidate)
+    return h, w
+
+
+def self_attn_forward(params: HypertileParams, scale_depth=True) -> Callable:
+
+    @wraps(params.forward)
+    def wrapper(*args, **kwargs):
+        if not params.enabled:
+            return params.forward(*args, **kwargs)
+
+        latent_tile_size = max(128, params.tile_size) // 8
+        x = args[0]
+
+        # VAE
+        if x.ndim == 4:
+            b, c, h, w = x.shape
+
+            nh = random_divisor(h, latent_tile_size, params.swap_size)
+            nw = random_divisor(w, latent_tile_size, params.swap_size)
+
+            if nh * nw > 1:
+                x = rearrange(x, "b c (nh h) (nw w) -> (b nh nw) c h w", nh=nh, nw=nw)  # split into nh * nw tiles
+
+            out = params.forward(x, *args[1:], **kwargs)
+
+            if nh * nw > 1:
+                out = rearrange(out, "(b nh nw) c h w -> b c (nh h) (nw w)", nh=nh, nw=nw)
+
+        # U-Net
+        else:
+            hw: int = x.size(1)
+            h, w = find_hw_candidates(hw, params.aspect_ratio)
+            assert h * w == hw, f"Invalid aspect ratio {params.aspect_ratio} for input of shape {x.shape}, hw={hw}, h={h}, w={w}"
+
+            factor = 2 ** params.depth if scale_depth else 1
+            nh = random_divisor(h, latent_tile_size * factor, params.swap_size)
+            nw = random_divisor(w, latent_tile_size * factor, params.swap_size)
+
+            if nh * nw > 1:
+                x = rearrange(x, "b (nh h nw w) c -> (b nh nw) (h w) c", h=h // nh, w=w // nw, nh=nh, nw=nw)
+
+            out = params.forward(x, *args[1:], **kwargs)
+
+            if nh * nw > 1:
+                out = rearrange(out, "(b nh nw) hw c -> b nh nw hw c", nh=nh, nw=nw)
+                out = rearrange(out, "b nh nw (h w) c -> b (nh h nw w) c", h=h // nh, w=w // nw)
+
+        return out
+
+    return wrapper
+
+
+def hypertile_hook_model(model: nn.Module, width, height, *, enable=False, tile_size_max=128, swap_size=1, max_depth=3, is_sdxl=False):
+    hypertile_layers = getattr(model, "__webui_hypertile_layers", None)
+    if hypertile_layers is None:
+        if not enable:
+            return
+
+        hypertile_layers = {}
+        layers = DEPTH_LAYERS_XL if is_sdxl else DEPTH_LAYERS
+
+        for depth in range(4):
+            for layer_name, module in model.named_modules():
+                if any(layer_name.endswith(try_name) for try_name in layers[depth]):
+                    params = HypertileParams()
+                    module.__webui_hypertile_params = params
+                    params.forward = module.forward
+                    params.depth = depth
+                    params.layer_name = layer_name
+                    module.forward = self_attn_forward(params)
+
+                    hypertile_layers[layer_name] = 1
+
+        model.__webui_hypertile_layers = hypertile_layers
+
+    aspect_ratio = width / height
+    tile_size = min(largest_tile_size_available(width, height), tile_size_max)
+
+    for layer_name, module in model.named_modules():
+        if layer_name in hypertile_layers:
+            params = module.__webui_hypertile_params
+
+            params.tile_size = tile_size
+            params.swap_size = swap_size
+            params.aspect_ratio = aspect_ratio
+            params.enabled = enable and params.depth <= max_depth
diff --git a/extensions-builtin/hypertile/scripts/hypertile_script.py b/extensions-builtin/hypertile/scripts/hypertile_script.py
new file mode 100644
index 00000000..395d584b
--- /dev/null
+++ b/extensions-builtin/hypertile/scripts/hypertile_script.py
@@ -0,0 +1,109 @@
+import hypertile

+from modules import scripts, script_callbacks, shared

+from scripts.hypertile_xyz import add_axis_options

+

+

+class ScriptHypertile(scripts.Script):

+    name = "Hypertile"

+

+    def title(self):

+        return self.name

+

+    def show(self, is_img2img):

+        return scripts.AlwaysVisible

+

+    def process(self, p, *args):

+        hypertile.set_hypertile_seed(p.all_seeds[0])

+

+        configure_hypertile(p.width, p.height, enable_unet=shared.opts.hypertile_enable_unet)

+

+        self.add_infotext(p)

+

+    def before_hr(self, p, *args):

+

+        enable = shared.opts.hypertile_enable_unet_secondpass or shared.opts.hypertile_enable_unet

+

+        # exclusive hypertile seed for the second pass

+        if enable:

+            hypertile.set_hypertile_seed(p.all_seeds[0])

+

+        configure_hypertile(p.hr_upscale_to_x, p.hr_upscale_to_y, enable_unet=enable)

+

+        if enable and not shared.opts.hypertile_enable_unet:

+            p.extra_generation_params["Hypertile U-Net second pass"] = True

+

+            self.add_infotext(p, add_unet_params=True)

+

+    def add_infotext(self, p, add_unet_params=False):

+        def option(name):

+            value = getattr(shared.opts, name)

+            default_value = shared.opts.get_default(name)

+            return None if value == default_value else value

+

+        if shared.opts.hypertile_enable_unet:

+            p.extra_generation_params["Hypertile U-Net"] = True

+

+        if shared.opts.hypertile_enable_unet or add_unet_params:

+            p.extra_generation_params["Hypertile U-Net max depth"] = option('hypertile_max_depth_unet')

+            p.extra_generation_params["Hypertile U-Net max tile size"] = option('hypertile_max_tile_unet')

+            p.extra_generation_params["Hypertile U-Net swap size"] = option('hypertile_swap_size_unet')

+

+        if shared.opts.hypertile_enable_vae:

+            p.extra_generation_params["Hypertile VAE"] = True

+            p.extra_generation_params["Hypertile VAE max depth"] = option('hypertile_max_depth_vae')

+            p.extra_generation_params["Hypertile VAE max tile size"] = option('hypertile_max_tile_vae')

+            p.extra_generation_params["Hypertile VAE swap size"] = option('hypertile_swap_size_vae')

+

+

+def configure_hypertile(width, height, enable_unet=True):

+    hypertile.hypertile_hook_model(

+        shared.sd_model.first_stage_model,

+        width,

+        height,

+        swap_size=shared.opts.hypertile_swap_size_vae,

+        max_depth=shared.opts.hypertile_max_depth_vae,

+        tile_size_max=shared.opts.hypertile_max_tile_vae,

+        enable=shared.opts.hypertile_enable_vae,

+    )

+

+    hypertile.hypertile_hook_model(

+        shared.sd_model.model,

+        width,

+        height,

+        swap_size=shared.opts.hypertile_swap_size_unet,

+        max_depth=shared.opts.hypertile_max_depth_unet,

+        tile_size_max=shared.opts.hypertile_max_tile_unet,

+        enable=enable_unet,

+        is_sdxl=shared.sd_model.is_sdxl

+    )

+

+

+def on_ui_settings():

+    import gradio as gr

+

+    options = {

+        "hypertile_explanation": shared.OptionHTML("""

+    <a href='https://github.com/tfernd/HyperTile'>Hypertile</a> optimizes the self-attention layer within U-Net and VAE models,

+    resulting in a reduction in computation time ranging from 1 to 4 times. The larger the generated image is, the greater the

+    benefit.

+    """),

+

+        "hypertile_enable_unet": shared.OptionInfo(False, "Enable Hypertile U-Net", infotext="Hypertile U-Net").info("enables hypertile for all modes, including hires fix second pass; noticeable change in details of the generated picture"),

+        "hypertile_enable_unet_secondpass": shared.OptionInfo(False, "Enable Hypertile U-Net for hires fix second pass", infotext="Hypertile U-Net second pass").info("enables hypertile just for hires fix second pass - regardless of whether the above setting is enabled"),

+        "hypertile_max_depth_unet": shared.OptionInfo(3, "Hypertile U-Net max depth", gr.Slider, {"minimum": 0, "maximum": 3, "step": 1}, infotext="Hypertile U-Net max depth").info("larger = more neural network layers affected; minor effect on performance"),

+        "hypertile_max_tile_unet": shared.OptionInfo(256, "Hypertile U-Net max tile size", gr.Slider, {"minimum": 0, "maximum": 512, "step": 16}, infotext="Hypertile U-Net max tile size").info("larger = worse performance"),

+        "hypertile_swap_size_unet": shared.OptionInfo(3, "Hypertile U-Net swap size", gr.Slider, {"minimum": 0, "maximum": 64, "step": 1}, infotext="Hypertile U-Net swap size"),

+

+        "hypertile_enable_vae": shared.OptionInfo(False, "Enable Hypertile VAE", infotext="Hypertile VAE").info("minimal change in the generated picture"),

+        "hypertile_max_depth_vae": shared.OptionInfo(3, "Hypertile VAE max depth", gr.Slider, {"minimum": 0, "maximum": 3, "step": 1}, infotext="Hypertile VAE max depth"),

+        "hypertile_max_tile_vae": shared.OptionInfo(128, "Hypertile VAE max tile size", gr.Slider, {"minimum": 0, "maximum": 512, "step": 16}, infotext="Hypertile VAE max tile size"),

+        "hypertile_swap_size_vae": shared.OptionInfo(3, "Hypertile VAE swap size ", gr.Slider, {"minimum": 0, "maximum": 64, "step": 1}, infotext="Hypertile VAE swap size"),

+    }

+

+    for name, opt in options.items():

+        opt.section = ('hypertile', "Hypertile")

+        shared.opts.add_option(name, opt)

+

+

+script_callbacks.on_ui_settings(on_ui_settings)

+script_callbacks.on_before_ui(add_axis_options)

diff --git a/extensions-builtin/hypertile/scripts/hypertile_xyz.py b/extensions-builtin/hypertile/scripts/hypertile_xyz.py
new file mode 100644
index 00000000..9e96ae3c
--- /dev/null
+++ b/extensions-builtin/hypertile/scripts/hypertile_xyz.py
@@ -0,0 +1,51 @@
+from modules import scripts
+from modules.shared import opts
+
+xyz_grid = [x for x in scripts.scripts_data if x.script_class.__module__ == "xyz_grid.py"][0].module
+
+def int_applier(value_name:str, min_range:int = -1, max_range:int = -1):
+    """
+    Returns a function that applies the given value to the given value_name in opts.data.
+    """
+    def validate(value_name:str, value:str):
+        value = int(value)
+        # validate value
+        if not min_range == -1:
+            assert value >= min_range, f"Value {value} for {value_name} must be greater than or equal to {min_range}"
+        if not max_range == -1:
+            assert value <= max_range, f"Value {value} for {value_name} must be less than or equal to {max_range}"
+    def apply_int(p, x, xs):
+        validate(value_name, x)
+        opts.data[value_name] = int(x)
+    return apply_int
+
+def bool_applier(value_name:str):
+    """
+    Returns a function that applies the given value to the given value_name in opts.data.
+    """
+    def validate(value_name:str, value:str):
+        assert value.lower() in ["true", "false"], f"Value {value} for {value_name} must be either true or false"
+    def apply_bool(p, x, xs):
+        validate(value_name, x)
+        value_boolean = x.lower() == "true"
+        opts.data[value_name] = value_boolean
+    return apply_bool
+
+def add_axis_options():
+    extra_axis_options = [
+        xyz_grid.AxisOption("[Hypertile] Unet First pass Enabled", str, bool_applier("hypertile_enable_unet"), choices=xyz_grid.boolean_choice(reverse=True)),
+        xyz_grid.AxisOption("[Hypertile] Unet Second pass Enabled", str, bool_applier("hypertile_enable_unet_secondpass"), choices=xyz_grid.boolean_choice(reverse=True)),
+        xyz_grid.AxisOption("[Hypertile] Unet Max Depth", int, int_applier("hypertile_max_depth_unet", 0, 3), choices=lambda: [str(x) for x in range(4)]),
+        xyz_grid.AxisOption("[Hypertile] Unet Max Tile Size", int, int_applier("hypertile_max_tile_unet", 0, 512)),
+        xyz_grid.AxisOption("[Hypertile] Unet Swap Size", int, int_applier("hypertile_swap_size_unet", 0, 64)),
+        xyz_grid.AxisOption("[Hypertile] VAE Enabled", str, bool_applier("hypertile_enable_vae"), choices=xyz_grid.boolean_choice(reverse=True)),
+        xyz_grid.AxisOption("[Hypertile] VAE Max Depth", int, int_applier("hypertile_max_depth_vae", 0, 3), choices=lambda: [str(x) for x in range(4)]),
+        xyz_grid.AxisOption("[Hypertile] VAE Max Tile Size", int, int_applier("hypertile_max_tile_vae", 0, 512)),
+        xyz_grid.AxisOption("[Hypertile] VAE Swap Size", int, int_applier("hypertile_swap_size_vae", 0, 64)),
+    ]
+    set_a = {opt.label for opt in xyz_grid.axis_options}
+    set_b = {opt.label for opt in extra_axis_options}
+    if set_a.intersection(set_b):
+        return
+
+    xyz_grid.axis_options.extend(extra_axis_options)