From 5a54039dd7b6482b99fde36add5753fd07f155ad Mon Sep 17 00:00:00 2001 From: StAlKeR7779 Date: Wed, 12 Apr 2023 22:54:25 +0300 Subject: [PATCH 1/8] Fix imports for diffusers 0.15.0 Imports: `diffusers.models.cross_attention -> diffusers.models.attention_processor` Unions: `AttnProcessor -> AttentionProcessor` Classes: | Old name | New name| | --- | --- | | CrossAttention | Attention | | CrossAttnProcessor | AttnProcessor | | XFormersCrossAttnProcessor | XFormersAttnProcessor | | CrossAttnAddedKVProcessor | AttnAddedKVProcessor | | LoRACrossAttnProcessor | LoRAAttnProcessor | | LoRAXFormersCrossAttnProcessor | LoRAXFormersAttnProcessor | Same names in this class: `SlicedAttnProcessor, SlicedAttnAddedKVProcessor` --- .../diffusion/cross_attention_control.py | 22 +++++++++---------- .../diffusion/shared_invokeai_diffusion.py | 6 ++--- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/ldm/models/diffusion/cross_attention_control.py b/ldm/models/diffusion/cross_attention_control.py index 7a66448b1cb..d73f3af260d 100644 --- a/ldm/models/diffusion/cross_attention_control.py +++ b/ldm/models/diffusion/cross_attention_control.py @@ -14,7 +14,7 @@ from compel.cross_attention_control import Arguments from diffusers.models.unet_2d_condition import UNet2DConditionModel -from diffusers.models.cross_attention import AttnProcessor +from diffusers.models.attention_processor import AttentionProcessor from ldm.invoke.devices import torch_dtype @@ -163,7 +163,7 @@ def offload_saved_attention_slices_to_cpu(self): class InvokeAICrossAttentionMixin: """ - Enable InvokeAI-flavoured CrossAttention calculation, which does aggressive low-memory slicing and calls + Enable InvokeAI-flavoured Attention calculation, which does aggressive low-memory slicing and calls through both to an attention_slice_wrangler and a slicing_strategy_getter for custom attention map wrangling and dymamic slicing strategy selection. """ @@ -178,7 +178,7 @@ def set_attention_slice_wrangler(self, wrangler: Optional[Callable[[nn.Module, t Set custom attention calculator to be called when attention is calculated :param wrangler: Callback, with args (module, suggested_attention_slice, dim, offset, slice_size), which returns either the suggested_attention_slice or an adjusted equivalent. - `module` is the current CrossAttention module for which the callback is being invoked. + `module` is the current Attention module for which the callback is being invoked. `suggested_attention_slice` is the default-calculated attention slice `dim` is -1 if the attenion map has not been sliced, or 0 or 1 for dimension-0 or dimension-1 slicing. If `dim` is >= 0, `offset` and `slice_size` specify the slice start and length. @@ -289,10 +289,10 @@ def get_invokeai_attention_mem_efficient(self, q, k, v): -def restore_default_cross_attention(model, is_running_diffusers: bool, processors_to_restore: Optional[AttnProcessor]=None): +def restore_default_cross_attention(model, is_running_diffusers: bool, processors_to_restore: Optional[AttentionProcessor]=None): if is_running_diffusers: unet = model - unet.set_attn_processor(processors_to_restore or CrossAttnProcessor()) + unet.set_attn_processor(processors_to_restore or AttnProcessor()) else: remove_attention_function(model) @@ -342,7 +342,7 @@ def override_cross_attention(model, context: Context, is_running_diffusers = Fal def get_cross_attention_modules(model, which: CrossAttentionType) -> list[tuple[str, InvokeAICrossAttentionMixin]]: - from ldm.modules.attention import CrossAttention # avoid circular import + from ldm.modules.attention import CrossAttention # avoid circular import # TODO: rename as in diffusers? cross_attention_class: type = InvokeAIDiffusersCrossAttention if isinstance(model,UNet2DConditionModel) else CrossAttention which_attn = "attn1" if which is CrossAttentionType.SELF else "attn2" attention_module_tuples = [(name,module) for name, module in model.named_modules() if @@ -448,7 +448,7 @@ def get_mem_free_total(device): -class InvokeAIDiffusersCrossAttention(diffusers.models.attention.CrossAttention, InvokeAICrossAttentionMixin): +class InvokeAIDiffusersCrossAttention(diffusers.models.attention.Attention, InvokeAICrossAttentionMixin): def __init__(self, **kwargs): super().__init__(**kwargs) @@ -473,8 +473,8 @@ def _attention(self, query, key, value, attention_mask=None): """ # base implementation -class CrossAttnProcessor: - def __call__(self, attn: CrossAttention, hidden_states, encoder_hidden_states=None, attention_mask=None): +class AttnProcessor: + def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, attention_mask=None): batch_size, sequence_length, _ = hidden_states.shape attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length) @@ -503,7 +503,7 @@ def __call__(self, attn: CrossAttention, hidden_states, encoder_hidden_states=No import torch -from diffusers.models.cross_attention import CrossAttention, CrossAttnProcessor, SlicedAttnProcessor +from diffusers.models.attention_processor import Attention, AttnProcessor, SlicedAttnProcessor @dataclass @@ -548,7 +548,7 @@ class SlicedSwapCrossAttnProcesser(SlicedAttnProcessor): # TODO: dynamically pick slice size based on memory conditions - def __call__(self, attn: CrossAttention, hidden_states, encoder_hidden_states=None, attention_mask=None, + def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, attention_mask=None, # kwargs swap_cross_attn_context: SwapCrossAttnContext=None): diff --git a/ldm/models/diffusion/shared_invokeai_diffusion.py b/ldm/models/diffusion/shared_invokeai_diffusion.py index f9518fca7ff..dcd29382393 100644 --- a/ldm/models/diffusion/shared_invokeai_diffusion.py +++ b/ldm/models/diffusion/shared_invokeai_diffusion.py @@ -5,7 +5,7 @@ import numpy as np import torch -from diffusers.models.cross_attention import AttnProcessor +from diffusers.models.attention_processor import AttentionProcessor from typing_extensions import TypeAlias from ldm.invoke.globals import Globals @@ -110,7 +110,7 @@ def custom_attention_context( def override_attention_processors( self, conditioning: ExtraConditioningInfo, step_count: int - ) -> Dict[str, AttnProcessor]: + ) -> Dict[str, AttentionProcessor]: """ setup cross attention .swap control. for diffusers this replaces the attention processor, so the previous attention processor is returned so that the caller can restore it later. @@ -135,7 +135,7 @@ def override_attention_processors( return old_attn_processors def restore_default_cross_attention( - self, processors_to_restore: Optional[dict[str, "AttnProcessor"]] = None + self, processors_to_restore: Optional[dict[str, "AttentionProcessor"]] = None ): self.cross_attention_control_context = None restore_default_cross_attention( From e24dd97b80fecc3c3183420197853ad6fbc46761 Mon Sep 17 00:00:00 2001 From: StAlKeR7779 Date: Wed, 12 Apr 2023 23:40:14 +0300 Subject: [PATCH 2/8] Fix that config attributes no longer accessible as object attributes --- ldm/invoke/generator/diffusers_pipeline.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ldm/invoke/generator/diffusers_pipeline.py b/ldm/invoke/generator/diffusers_pipeline.py index 2dd90c8d1e1..2202cec841a 100644 --- a/ldm/invoke/generator/diffusers_pipeline.py +++ b/ldm/invoke/generator/diffusers_pipeline.py @@ -400,7 +400,7 @@ def device(self) -> torch.device: @property def _submodels(self) -> Sequence[torch.nn.Module]: module_names, _, _ = self.extract_init_dict(dict(self.config)) - values = [getattr(self, name) for name in module_names.keys()] + values = [getattr(self, name, getattr(self.config, name)) for name in module_names.keys()] return [m for m in values if isinstance(m, torch.nn.Module)] def image_from_embeddings(self, latents: torch.Tensor, num_inference_steps: int, From 16c97ca0cb02937e6128a61d016f2e6518fe5038 Mon Sep 17 00:00:00 2001 From: StAlKeR7779 Date: Wed, 12 Apr 2023 23:57:45 +0300 Subject: [PATCH 3/8] Fix num_train_timesteps in config --- ldm/invoke/generator/diffusers_pipeline.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ldm/invoke/generator/diffusers_pipeline.py b/ldm/invoke/generator/diffusers_pipeline.py index 2202cec841a..003eba21dee 100644 --- a/ldm/invoke/generator/diffusers_pipeline.py +++ b/ldm/invoke/generator/diffusers_pipeline.py @@ -471,7 +471,7 @@ def generate_latents_from_embeddings(self, latents: torch.Tensor, timesteps, step_count=len(self.scheduler.timesteps) ): - yield PipelineIntermediateState(run_id=run_id, step=-1, timestep=self.scheduler.num_train_timesteps, + yield PipelineIntermediateState(run_id=run_id, step=-1, timestep=self.scheduler.config.num_train_timesteps, latents=latents) batch_size = latents.shape[0] From 0bc5dcc6634ff746b0e0ed040932a978b8ba0f91 Mon Sep 17 00:00:00 2001 From: StAlKeR7779 Date: Thu, 13 Apr 2023 16:05:04 +0300 Subject: [PATCH 4/8] Refactor --- ldm/invoke/generator/diffusers_pipeline.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/ldm/invoke/generator/diffusers_pipeline.py b/ldm/invoke/generator/diffusers_pipeline.py index 003eba21dee..dacd2797649 100644 --- a/ldm/invoke/generator/diffusers_pipeline.py +++ b/ldm/invoke/generator/diffusers_pipeline.py @@ -400,8 +400,15 @@ def device(self) -> torch.device: @property def _submodels(self) -> Sequence[torch.nn.Module]: module_names, _, _ = self.extract_init_dict(dict(self.config)) - values = [getattr(self, name, getattr(self.config, name)) for name in module_names.keys()] - return [m for m in values if isinstance(m, torch.nn.Module)] + submodels = [] + for name in module_names.keys(): + if hasattr(self, name): + value = getattr(self, name) + else: + value = getattr(self.config, name) + if isinstance(value, torch.nn.Module): + submodels.append(value) + return submodels def image_from_embeddings(self, latents: torch.Tensor, num_inference_steps: int, conditioning_data: ConditioningData, From c6ecf3afc5fab9213ecf96fe01d6b4b077362273 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Thu, 13 Apr 2023 22:38:50 -0400 Subject: [PATCH 5/8] pin diffusers to 0.15.*, and fix deprecation warning on unet.in_channels --- ldm/invoke/generator/diffusers_pipeline.py | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/ldm/invoke/generator/diffusers_pipeline.py b/ldm/invoke/generator/diffusers_pipeline.py index dacd2797649..fc764726f71 100644 --- a/ldm/invoke/generator/diffusers_pipeline.py +++ b/ldm/invoke/generator/diffusers_pipeline.py @@ -762,7 +762,7 @@ def _tokenize(self, prompt: Union[str, List[str]]): @property def channels(self) -> int: """Compatible with DiffusionWrapper""" - return self.unet.in_channels + return self.unet.config.in_channels def decode_latents(self, latents): # Explicit call to get the vae loaded, since `decode` isn't the forward method. diff --git a/pyproject.toml b/pyproject.toml index 1b22eaca4de..73c80d0d2b4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,7 +34,7 @@ dependencies = [ "clip_anytorch", "compel~=1.1.0", "datasets", - "diffusers[torch]~=0.14", + "diffusers[torch]==0.15.*", "dnspython==2.2.1", "einops", "eventlet", From 2a2c86896ab89d83f53e77b7b2dffeb1eeb07981 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Thu, 20 Apr 2023 13:29:20 -0400 Subject: [PATCH 6/8] pull in diffusers 0.15.1 - Change diffusers dependency to `diffusers~=0.15.0` which *should* enforce non-breaking changes. --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 73c80d0d2b4..f0b0cb6b6c5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,7 +34,7 @@ dependencies = [ "clip_anytorch", "compel~=1.1.0", "datasets", - "diffusers[torch]==0.15.*", + "diffusers[torch]~=0.15.0", "dnspython==2.2.1", "einops", "eventlet", From 830740b93b139f552db24e858c20509a7dbaa57b Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Tue, 25 Apr 2023 07:05:07 -0400 Subject: [PATCH 7/8] remove redundant/buggy restore_default_attention() method --- ldm/models/diffusion/shared_invokeai_diffusion.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/ldm/models/diffusion/shared_invokeai_diffusion.py b/ldm/models/diffusion/shared_invokeai_diffusion.py index d4338879dca..853662ce4ea 100644 --- a/ldm/models/diffusion/shared_invokeai_diffusion.py +++ b/ldm/models/diffusion/shared_invokeai_diffusion.py @@ -149,16 +149,6 @@ def override_attention_processors( ) return old_attn_processors - def restore_default_cross_attention( - self, processors_to_restore: Optional[dict[str, "AttentionProcessor"]] = None - ): - self.cross_attention_control_context = None - restore_default_cross_attention( - self.model, - is_running_diffusers=self.is_running_diffusers, - processors_to_restore=processors_to_restore, - ) - def setup_attention_map_saving(self, saver: AttentionMapSaver): def callback(slice, dim, offset, slice_size, key): if dim is not None: From 432e526999cea8c291825423aa87ad63b19d9757 Mon Sep 17 00:00:00 2001 From: Sergey Borisov Date: Tue, 25 Apr 2023 14:49:08 +0300 Subject: [PATCH 8/8] Revert merge changes --- .../diffusion/cross_attention_control.py | 8 ----- .../diffusion/shared_invokeai_diffusion.py | 29 +------------------ 2 files changed, 1 insertion(+), 36 deletions(-) diff --git a/ldm/models/diffusion/cross_attention_control.py b/ldm/models/diffusion/cross_attention_control.py index 6f963b031a1..8c0e495df5b 100644 --- a/ldm/models/diffusion/cross_attention_control.py +++ b/ldm/models/diffusion/cross_attention_control.py @@ -14,7 +14,6 @@ from compel.cross_attention_control import Arguments from diffusers.models.unet_2d_condition import UNet2DConditionModel -from diffusers.models.attention_processor import AttentionProcessor from ldm.invoke.devices import torch_dtype @@ -288,13 +287,6 @@ def get_invokeai_attention_mem_efficient(self, q, k, v): return self.einsum_op_tensor_mem(q, k, v, 32) -def restore_default_cross_attention(model, is_running_diffusers: bool, processors_to_restore: Optional[AttentionProcessor]=None): - if is_running_diffusers: - unet = model - unet.set_attn_processor(processors_to_restore or AttnProcessor()) - else: - remove_attention_function(model) - def setup_cross_attention_control_attention_processors(unet: UNet2DConditionModel, context: Context): """ Inject attention parameters and functions into the passed in model to enable cross attention editing. diff --git a/ldm/models/diffusion/shared_invokeai_diffusion.py b/ldm/models/diffusion/shared_invokeai_diffusion.py index 853662ce4ea..61a78147160 100644 --- a/ldm/models/diffusion/shared_invokeai_diffusion.py +++ b/ldm/models/diffusion/shared_invokeai_diffusion.py @@ -1,11 +1,10 @@ from contextlib import contextmanager from dataclasses import dataclass from math import ceil -from typing import Callable, Optional, Union, Any, Dict +from typing import Callable, Optional, Union, Any import numpy as np import torch -from diffusers.models.attention_processor import AttentionProcessor from diffusers import UNet2DConditionModel from typing_extensions import TypeAlias @@ -123,32 +122,6 @@ def custom_attention_context( # TODO resuscitate attention map saving # self.remove_attention_map_saving() - def override_attention_processors( - self, conditioning: ExtraConditioningInfo, step_count: int - ) -> Dict[str, AttentionProcessor]: - """ - setup cross attention .swap control. for diffusers this replaces the attention processor, so - the previous attention processor is returned so that the caller can restore it later. - """ - old_attn_processors = self.model.attn_processors - - # Load lora conditions into the model - if conditioning.has_lora_conditions: - for condition in conditioning.lora_conditions: - condition(self.model) - - if conditioning.wants_cross_attention_control: - self.cross_attention_control_context = Context( - arguments=conditioning.cross_attention_control_args, - step_count=step_count, - ) - override_cross_attention( - self.model, - self.cross_attention_control_context, - is_running_diffusers=self.is_running_diffusers, - ) - return old_attn_processors - def setup_attention_map_saving(self, saver: AttentionMapSaver): def callback(slice, dim, offset, slice_size, key): if dim is not None: