add linfusion

Signed-off-by: Vladimir Mandic <mandic00@live.com>

modules/linfusion/__init__.py

@@ -0,0 +1,44 @@
+from modules import shared, sd_models, devices
+from .linfusion import LinFusion
+from .attention import GeneralizedLinearAttention
+
+
+applied: LinFusion = None
+
+
+def detect(pipeline):
+    if pipeline.__class__.__name__ == 'StableDiffusionXLPipeline':
+        return "Yuanshi/LinFusion-XL"
+    if pipeline.__class__.__name__ == 'StableDiffusionPipeline':
+        return "Yuanshi/LinFusion-1-5"
+    return None
+
+
+def apply(pipeline, pretrained: bool = True):
+    global applied # pylint: disable=global-statement
+    if applied is not None:
+        return
+    # linfusion = LinFusion.construct_for(pipeline=pipeline)
+    if not pretrained:
+        model_path = None
+        default_config = LinFusion.get_default_config(unet=pipeline.unet)
+        applied = LinFusion(**default_config).to(device=pipeline.unet.device, dtype=pipeline.unet.dtype)
+        applied.mount_to(unet=pipeline.unet)
+    else:
+        model_path = detect(pipeline)
+        if model_path is None:
+            shared.log.error('LinFusion: unsupported model type')
+            return
+        applied = LinFusion.from_pretrained(model_path, cache_dir=shared.opts.hfcache_dir).to(device=pipeline.unet.device, dtype=pipeline.unet.dtype)
+        applied.mount_to(unet=pipeline.unet)
+    shared.log.debug(f'LinFusion: apply class={applied.__class__.__name__} model="{model_path}" modules={len(applied.modules_dict)}')
+
+
+def unapply(pipeline):
+    global applied # pylint: disable=global-statement
+    if applied is None:
+        return
+    shared.log.debug('LinFusion: unapply')
+    sd_models.set_diffusers_attention(pipeline)
+    devices.torch_gc()
+    applied = None

modules/linfusion/attention.py

@@ -0,0 +1,83 @@
+import torch
+import torch.nn.functional as F
+from diffusers.models.attention_processor import Attention
+
+
+def get_none_linear_projection(query_dim, mid_dim=None):
+    # If mid_dim is None, then the mid_dim is the same as query_dim
+    # If mid_dim is -1, then no non-linear projection is used, and the identity is returned
+    return (
+        torch.nn.Sequential(
+            torch.nn.Linear(query_dim, mid_dim or query_dim),
+            torch.nn.LayerNorm(mid_dim or query_dim),
+            torch.nn.LeakyReLU(inplace=True),
+            torch.nn.Linear(mid_dim or query_dim, query_dim),
+        )
+        if mid_dim != -1
+        else torch.nn.Identity()
+    )
+
+
+class GeneralizedLinearAttention(Attention):
+    def __init__(self, *args, projection_mid_dim=None, **kwargs):
+        """
+        Args:
+            query_dim: the dimension of the query.
+            out_dim: the dimension of the output.
+            dim_head: the dimension of the head. (dim_head * num_heads = query_dim)
+            projection_mid_dim: the dimension of the intermediate layer in the non-linear projection.
+              If `None`, then the dimension is the same as the query dimension.
+              If `-1`, then no non-linear projection is used, and the identity is returned.
+        """
+        super().__init__(*args, **kwargs)
+        self.add_non_linear_model(projection_mid_dim)
+
+    def from_attention_instance(self, attention_instance, projection_mid_dim=None):
+        assert isinstance(attention_instance, Attention)
+        new_instance = GeneralizedLinearAttention(128)
+        new_instance.__dict__ = attention_instance.__dict__
+        new_instance.add_non_linear_model(mid_dim = projection_mid_dim)
+        return new_instance
+
+    def add_non_linear_model(self, mid_dim=None, **kwargs):
+        query_dim = self.to_q.weight.shape[0]
+        self.to_q_ = get_none_linear_projection(query_dim, mid_dim, **kwargs)
+        self.to_k_ = get_none_linear_projection(query_dim, mid_dim, **kwargs)
+
+    def forward(
+        self,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        **kwargs,
+    ):
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+
+        _, sequence_length, _ = hidden_states.shape
+
+        query = self.to_q(hidden_states + self.to_q_(hidden_states))
+        key = self.to_k(encoder_hidden_states + self.to_k_(encoder_hidden_states))
+        value = self.to_v(encoder_hidden_states)
+
+        query = self.head_to_batch_dim(query)
+        key = self.head_to_batch_dim(key)
+        value = self.head_to_batch_dim(value)
+
+        query = F.elu(query) + 1.0
+        key = F.elu(key) + 1.0
+
+        z = query @ key.mean(dim=-2, keepdim=True).transpose(-2, -1) + 1e-4
+        kv = (key.transpose(-2, -1) * (sequence_length**-0.5)) @ (
+            value * (sequence_length**-0.5)
+        )
+        hidden_states = query @ kv / z
+
+        hidden_states = self.batch_to_head_dim(hidden_states)
+
+        # linear proj
+        hidden_states = self.to_out[0](hidden_states)
+        # dropout
+        hidden_states = self.to_out[1](hidden_states)
+
+        return hidden_states

modules/linfusion/linfusion.py

@@ -0,0 +1,119 @@
+import functools
+from diffusers.models.attention_processor import Attention
+from diffusers import ModelMixin, ConfigMixin
+from .attention import GeneralizedLinearAttention
+
+
+model_dict = {
+    "runwayml/stable-diffusion-v1-5": "Yuanshi/LinFusion-1-5",
+    "SG161222/Realistic_Vision_V4.0_noVAE": "Yuanshi/LinFusion-1-5",
+    "Lykon/dreamshaper-8": "Yuanshi/LinFusion-1-5",
+    "stabilityai/stable-diffusion-2-1": "Yuanshi/LinFusion-2-1",
+    "stabilityai/stable-diffusion-xl-base-1.0": "Yuanshi/LinFusion-XL",
+}
+
+
+def replace_submodule(model, module_name, new_submodule):
+    path, attr = module_name.rsplit(".", 1)
+    parent_module = functools.reduce(getattr, path.split("."), model)
+    setattr(parent_module, attr, new_submodule)
+
+
+class LinFusion(ModelMixin, ConfigMixin):
+    def __init__(self, modules_list, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+        self.modules_dict = {}
+        self.register_to_config(modules_list=modules_list)
+
+        for i, attention_config in enumerate(modules_list):
+            dim_n = attention_config["dim_n"]
+            heads = attention_config["heads"]
+            projection_mid_dim = attention_config["projection_mid_dim"]
+            linear_attention = GeneralizedLinearAttention(
+                query_dim=dim_n,
+                out_dim=dim_n,
+                dim_head=dim_n // heads,
+                projection_mid_dim=projection_mid_dim,
+            )
+            self.add_module(f"{i}", linear_attention)
+            self.modules_dict[attention_config["module_name"]] = linear_attention
+
+    @classmethod
+    def get_default_config(
+        cls,
+        pipeline=None,
+        unet=None,
+    ):
+        """
+        Get the default configuration for the LinFusion model.
+        (The `projection_mid_dim` is same as the `query_dim` by default.)
+        """
+        assert unet is not None or pipeline.unet is not None
+        unet = unet or pipeline.unet
+        modules_list = []
+        for module_name, module in unet.named_modules():
+            if not isinstance(module, Attention):
+                continue
+            if "attn1" not in module_name:
+                continue
+            dim_n = module.to_q.weight.shape[0]
+            # modules_list.append((module_name, dim_n, module.heads))
+            modules_list.append(
+                {
+                    "module_name": module_name,
+                    "dim_n": dim_n,
+                    "heads": module.heads,
+                    "projection_mid_dim": None,
+                }
+            )
+        return {"modules_list": modules_list}
+
+    @classmethod
+    def construct_for(
+        cls,
+        pipeline=None,
+        unet=None,
+        load_pretrained=True,
+        pretrained_model_name_or_path=None,
+        pipe_name_path=None,
+    ) -> "LinFusion":
+        """
+        Construct a LinFusion object for the given pipeline.
+        """
+        assert unet is not None or pipeline.unet is not None
+        unet = unet or pipeline.unet
+        if load_pretrained:
+            # Load from pretrained
+            if not pretrained_model_name_or_path:
+                pipe_name_path = pipe_name_path or pipeline._internal_dict._name_or_path
+                pretrained_model_name_or_path = model_dict.get(pipe_name_path, None)
+                if pretrained_model_name_or_path:
+                    print(
+                        f"Matching LinFusion '{pretrained_model_name_or_path}' for pipeline '{pipe_name_path}'."
+                    )
+                else:
+                    raise RuntimeError(
+                        f"LinFusion not found for pipeline [{pipe_name_path}], please provide the path."
+                    )
+            linfusion = (
+                LinFusion.from_pretrained(pretrained_model_name_or_path)
+                .to(unet.device)
+                .to(unet.dtype)
+            )
+        else:
+            # Create from scratch without pretrained parameters
+            default_config = LinFusion.get_default_config(unet=unet)
+            linfusion = LinFusion(**default_config).to(unet.device).to(unet.dtype)
+        linfusion.mount_to(unet=unet)
+        return linfusion
+
+    def mount_to(self, pipeline=None, unet=None) -> None:
+        """
+        Mounts the modules in the `modules_dict` to the given `pipeline`.
+        """
+        assert unet is not None or pipeline.unet is not None
+        unet = unet or pipeline.unet
+        for module_name, module in self.modules_dict.items():
+            replace_submodule(unet, module_name, module)
+        self.to(unet.device).to(unet.dtype)

modules/sd_models.py

@@ -1346,6 +1346,11 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
             sd_model = sd_models_compile.compile_diffusers(sd_model)
         timer.record("compile")
 
+        if shared.opts.enable_linfusion:
+            from modules import linfusion
+            linfusion.apply(sd_model)
+            timer.record("linfusion")
+
     except Exception as e:
         shared.log.error(f"Load {op}: {e}")
         errors.display(e, "Model")

modules/shared.py

@@ -570,6 +570,7 @@ options_templates.update(options_section(('diffusers', "Diffusers Settings"), {
     "diffusers_pooled": OptionInfo("default", "Diffusers SDXL pooled embeds", gr.Radio, {"choices": ['default', 'weighted']}),
     "diffusers_zeros_prompt_pad": OptionInfo(False, "Use zeros for prompt padding", gr.Checkbox),
     "huggingface_token": OptionInfo('', 'HuggingFace token'),
+    "enable_linfusion": OptionInfo(False, "Apply LinFusion distillation on load"),
 
     "onnx_sep": OptionInfo("<h2>ONNX Runtime</h2>", "", gr.HTML),
     "onnx_execution_provider": OptionInfo(execution_providers.get_default_execution_provider().value, 'Execution Provider', gr.Dropdown, lambda: {"choices": execution_providers.available_execution_providers }),