From bbaec297207b74a96d6a3bc1d5741e22f444bf6e Mon Sep 17 00:00:00 2001 From: Disty0 Date: Fri, 21 Jun 2024 21:07:15 +0300 Subject: [PATCH] Cleanup --- modules/intel/ipex/diffusers.py | 19 ++++++++----------- modules/sd_hijack_dynamic_atten.py | 27 ++++++++++++--------------- modules/sd_models.py | 6 +++--- 3 files changed, 23 insertions(+), 29 deletions(-) diff --git a/modules/intel/ipex/diffusers.py b/modules/intel/ipex/diffusers.py index 732a18568..4f294ce6b 100644 --- a/modules/intel/ipex/diffusers.py +++ b/modules/intel/ipex/diffusers.py @@ -70,8 +70,8 @@ class SlicedAttnProcessor: # pylint: disable=too-few-public-methods def __init__(self, slice_size): self.slice_size = slice_size - def __call__(self, attn: Attention, hidden_states: torch.FloatTensor, - encoder_hidden_states=None, attention_mask=None) -> torch.FloatTensor: # pylint: disable=too-many-statements, too-many-locals, too-many-branches + def __call__(self, attn: Attention, hidden_states: torch.Tensor, + encoder_hidden_states=None, attention_mask=None) -> torch.Tensor: # pylint: disable=too-many-statements, too-many-locals, too-many-branches residual = hidden_states @@ -188,14 +188,11 @@ class AttnProcessor: Default processor for performing attention-related computations. """ - def __call__(self, attn: Attention, hidden_states: torch.FloatTensor, - encoder_hidden_states=None, attention_mask=None, - temb=None, scale: float = 1.0) -> torch.Tensor: # pylint: disable=too-many-statements, too-many-locals, too-many-branches + def __call__(self, attn, hidden_states: torch.Tensor, encoder_hidden_states=None, attention_mask=None, + temb=None, *args, **kwargs) -> torch.Tensor: # pylint: disable=too-many-statements, too-many-locals, too-many-branches residual = hidden_states - args = () if USE_PEFT_BACKEND else (scale,) - if attn.spatial_norm is not None: hidden_states = attn.spatial_norm(hidden_states, temb) @@ -213,15 +210,15 @@ class AttnProcessor: if attn.group_norm is not None: hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - query = attn.to_q(hidden_states, *args) + query = attn.to_q(hidden_states) if encoder_hidden_states is None: encoder_hidden_states = hidden_states elif attn.norm_cross: encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - key = attn.to_k(encoder_hidden_states, *args) - value = attn.to_v(encoder_hidden_states, *args) + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) query = attn.head_to_batch_dim(query) key = attn.head_to_batch_dim(key) @@ -292,7 +289,7 @@ class AttnProcessor: hidden_states = attn.batch_to_head_dim(hidden_states) # linear proj - hidden_states = attn.to_out[0](hidden_states, *args) + hidden_states = attn.to_out[0](hidden_states) # dropout hidden_states = attn.to_out[1](hidden_states) diff --git a/modules/sd_hijack_dynamic_atten.py b/modules/sd_hijack_dynamic_atten.py index 5ba9dc6e4..b2d6fdc42 100644 --- a/modules/sd_hijack_dynamic_atten.py +++ b/modules/sd_hijack_dynamic_atten.py @@ -110,8 +110,8 @@ class DynamicAttnProcessorSDP: if not hasattr(F, "scaled_dot_product_attention"): raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") - def __call__( - self, attn, hidden_states: torch.FloatTensor, encoder_hidden_states=None, attention_mask=None, temb=None, scale: float = 1.0) -> torch.FloatTensor: + def __call__(self, attn, hidden_states: torch.Tensor, encoder_hidden_states=None, attention_mask=None, temb=None, *args, **kwargs) -> torch.Tensor: + residual = hidden_states if attn.spatial_norm is not None: hidden_states = attn.spatial_norm(hidden_states, temb) @@ -135,16 +135,15 @@ class DynamicAttnProcessorSDP: if attn.group_norm is not None: hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - args = () if USE_PEFT_BACKEND else (scale,) - query = attn.to_q(hidden_states, *args) + query = attn.to_q(hidden_states) if encoder_hidden_states is None: encoder_hidden_states = hidden_states elif attn.norm_cross: encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - key = attn.to_k(encoder_hidden_states, *args) - value = attn.to_v(encoder_hidden_states, *args) + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) inner_dim = key.shape[-1] head_dim = inner_dim // attn.heads @@ -167,7 +166,7 @@ class DynamicAttnProcessorSDP: hidden_states = hidden_states.to(query.dtype) # linear proj - hidden_states = attn.to_out[0](hidden_states, *args) + hidden_states = attn.to_out[0](hidden_states) # dropout hidden_states = attn.to_out[1](hidden_states) @@ -190,13 +189,11 @@ class DynamicAttnProcessorBMM: based on AttnProcessor V1 """ - def __call__(self, attn, hidden_states: torch.FloatTensor, encoder_hidden_states=None, attention_mask=None, - temb=None, scale: float = 1.0) -> torch.Tensor: # pylint: disable=too-many-statements, too-many-locals, too-many-branches + def __call__(self, attn, hidden_states: torch.Tensor, encoder_hidden_states=None, attention_mask=None, + temb=None, *args, **kwargs) -> torch.Tensor: # pylint: disable=too-many-statements, too-many-locals, too-many-branches residual = hidden_states - args = () if USE_PEFT_BACKEND else (scale,) - if attn.spatial_norm is not None: hidden_states = attn.spatial_norm(hidden_states, temb) @@ -214,15 +211,15 @@ class DynamicAttnProcessorBMM: if attn.group_norm is not None: hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - query = attn.to_q(hidden_states, *args) + query = attn.to_q(hidden_states) if encoder_hidden_states is None: encoder_hidden_states = hidden_states elif attn.norm_cross: encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - key = attn.to_k(encoder_hidden_states, *args) - value = attn.to_v(encoder_hidden_states, *args) + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) query = attn.head_to_batch_dim(query) key = attn.head_to_batch_dim(key) @@ -294,7 +291,7 @@ class DynamicAttnProcessorBMM: hidden_states = attn.batch_to_head_dim(hidden_states) # linear proj - hidden_states = attn.to_out[0](hidden_states, *args) + hidden_states = attn.to_out[0](hidden_states) # dropout hidden_states = attn.to_out[1](hidden_states) diff --git a/modules/sd_models.py b/modules/sd_models.py index 4aa561437..6efd56762 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -1299,9 +1299,6 @@ def switch_pipe(cls: diffusers.DiffusionPipeline, pipeline: diffusers.DiffusionP def set_diffuser_pipe(pipe, new_pipe_type): - if get_diffusers_task(pipe) == new_pipe_type: - return pipe - n = getattr(pipe.__class__, '__name__', '') if new_pipe_type == DiffusersTaskType.TEXT_2_IMAGE and 'StableDiffusionXL' in n and 'requires_aesthetics_score' in pipe.config and hasattr(pipe, '_internal_dict'): # diffusers adds requires_aesthetics_score with img2img and complains if requires_aesthetics_score exist in txt2img @@ -1310,6 +1307,9 @@ def set_diffuser_pipe(pipe, new_pipe_type): del pipe._internal_dict pipe.register_to_config(**internal_dict) + if get_diffusers_task(pipe) == new_pipe_type: + return pipe + # skip specific pipelines if n in ['StableDiffusionReferencePipeline', 'StableDiffusionAdapterPipeline', 'AnimateDiffPipeline', 'AnimateDiffSDXLPipeline']: return pipe