add linfusion

Signed-off-by: Vladimir Mandic <mandic00@live.com>
2026-09-17 08:19:11 +02:00 · 2024-10-09 11:01:16 -04:00
parent f17fbe096d
commit 723ce74986
9 changed files with 270 additions and 12 deletions
@@ -1,23 +1,24 @@
 # Change Log for SD.Next

-## Update for 2024-10-08
+## Update for 2024-10-09

-### Highlights for 2024-10-08
+### Highlights for 2024-10-09

 - **Reprocess**: New workflow options that allow you to generate at lower quality and then  
  reprocess at higher quality for select images only or generate without hires/refine and then reprocess with hires/refine  
  and you can pick any previous latent from auto-captured history!  
 - **Detailer** Fully built-in detailer workflow without with support for all standard models  
 - New fine-tuned [CLiP-ViT-L]((https://huggingface.co/zer0int/CLIP-GmP-ViT-L-14)) 1st stage **text-encoders** used by SD15, SDXL, Flux.1, etc. brings additional details to your images  
- Integration with [Ctrl+X](https://github.com/genforce/ctrl-x) which allows for control of **structure and appearance** without the need for extra models and  
-  [APG: Adaptive Projected Guidance](https://arxiv.org/pdf/2410.02416) for optimal **guidance** control  
+- Integration with [Ctrl+X](https://github.com/genforce/ctrl-x) which allows for control of **structure and appearance** without the need for extra models,  
+  [APG: Adaptive Projected Guidance](https://arxiv.org/pdf/2410.02416) for optimal **guidance** control,  
+  [LinFusion](https://github.com/Huage001/LinFusion) for on-the-fly distillation of any sd15/sdxl model  
 - Auto-detection of best available **device/dtype** settings for your platform and GPU reduces neeed for manual configuration  
 - Full rewrite of **sampler options**, not far more streamlined with tons of new options to tweak scheduler behavior  
 - Improved **LoRA** detection and handling for all supported models  

 And other goodies like multiple *XYZ grid* improvements, additional *Flux ControlNets*, additional *Interrogate models*, better *LoRA tags* support, and more...

-### Details for 2024-10-07
+### Details for 2024-10-09

 - **reprocess**
  - new top-level button: reprocess latent from your history of generated image(s)  
@@ -109,6 +110,11 @@ And other goodies like multiple *XYZ grid* improvements, additional *Flux Contro
  - for normal cfg scale, use negative momentum: e.g. cfg=6 => momentum=-0.3
  - for high   cfg scale, use neutral  momentum: e.g. cfg=10 => momentum=0.0

+- [LinFusion](https://github.com/Huage001/LinFusion)  
+  - apply liner distillation to during load to any sd15/sdxl model  
+  - can reduce vram use for high resolutions and increase performance
+  - *note*: use lower cfg scales as typical for distilled models  
+
 - **flux**  
  - avoid unet load if unchanged  
  - mark specific unet as unavailable if load failed  
@@ -449,7 +449,7 @@ def check_python(supported_minors=[9, 10, 11, 12], reason=None):

 # check diffusers version
 def check_diffusers():
-    sha = '33fafe3d143ca8380a9e405e7acfa69091d863fb'
+    sha = '31058cdaef63ca660a1a045281d156239fba8192'
    pkg = pkg_resources.working_set.by_key.get('diffusers', None)
    minor = int(pkg.version.split('.')[1] if pkg is not None else 0)
    cur = opts.get('diffusers_version', '') if minor > 0 else ''
@@ -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
@@ -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
@@ -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)
@@ -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")
@@ -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 }),
@@ -31,7 +31,7 @@ invisible-watermark
 pi-heif
 safetensors==0.4.5
 tensordict==0.1.2
-peft==0.11.1
+peft==0.13.1
 httpx==0.24.1
 compel==2.0.3
 torchsde==0.2.6
@@ -40,11 +40,11 @@ clip-interrogator==0.6.0
 antlr4-python3-runtime==4.9.3
 requests==2.32.3
 tqdm==4.66.5
-accelerate==0.33.0
+accelerate==1.0.0
 opencv-contrib-python-headless==4.9.0.80
 einops==0.4.1
 gradio==3.43.2
-huggingface_hub==0.25.1
+huggingface_hub==0.25.2
 numexpr==2.8.8
 numpy==1.26.4
 numba==0.59.1
@@ -54,13 +54,13 @@ pandas
 protobuf==4.25.3
 pytorch_lightning==1.9.4
 tokenizers==0.20.0
-transformers==4.45.1
+transformers==4.45.2
 urllib3==1.26.19
 Pillow==10.4.0
 timm==0.9.16
 pydantic==1.10.15
 pyparsing==3.1.4
-typing-extensions==4.11.0
+typing-extensions==4.12.2
 torchdiffeq
 dctorch
 scikit-image
@@ -155,7 +155,7 @@ def initialize():


 def load_model():
-    if not shared.opts.sd_checkpoint_autoload and shared.cmd_opts.ckpt is not None:
+    if not shared.opts.sd_checkpoint_autoload and shared.cmd_opts.ckpt is None:
        log.debug('Model auto load disabled')
    else:
        shared.state.begin('Load')