From 723ce749860ee78d2dcd2e574e067fa00057b488 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 9 Oct 2024 11:01:16 -0400 Subject: [PATCH] add linfusion Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 16 +++-- installer.py | 2 +- modules/linfusion/__init__.py | 44 ++++++++++++ modules/linfusion/attention.py | 83 +++++++++++++++++++++++ modules/linfusion/linfusion.py | 119 +++++++++++++++++++++++++++++++++ modules/sd_models.py | 5 ++ modules/shared.py | 1 + requirements.txt | 10 +-- webui.py | 2 +- 9 files changed, 270 insertions(+), 12 deletions(-) create mode 100644 modules/linfusion/__init__.py create mode 100644 modules/linfusion/attention.py create mode 100644 modules/linfusion/linfusion.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 041603411..0c82db9c9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/installer.py b/installer.py index 2fd19fc19..4c51a0bfa 100644 --- a/installer.py +++ b/installer.py @@ -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 '' diff --git a/modules/linfusion/__init__.py b/modules/linfusion/__init__.py new file mode 100644 index 000000000..b263002bb --- /dev/null +++ b/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 diff --git a/modules/linfusion/attention.py b/modules/linfusion/attention.py new file mode 100644 index 000000000..0437e7d5e --- /dev/null +++ b/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 diff --git a/modules/linfusion/linfusion.py b/modules/linfusion/linfusion.py new file mode 100644 index 000000000..908e7a08d --- /dev/null +++ b/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) diff --git a/modules/sd_models.py b/modules/sd_models.py index 1f9d37176..09fa1dc10 100644 --- a/modules/sd_models.py +++ b/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") diff --git a/modules/shared.py b/modules/shared.py index 91bceba38..b972a05b2 100644 --- a/modules/shared.py +++ b/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("

ONNX Runtime

", "", gr.HTML), "onnx_execution_provider": OptionInfo(execution_providers.get_default_execution_provider().value, 'Execution Provider', gr.Dropdown, lambda: {"choices": execution_providers.available_execution_providers }), diff --git a/requirements.txt b/requirements.txt index 5a546acd7..4999331ca 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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 diff --git a/webui.py b/webui.py index e2e7285db..cf244fc44 100644 --- a/webui.py +++ b/webui.py @@ -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')