From 9d17cf4c122b98b25b8cb9e3388c1a75df68cdb2 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Mon, 7 Aug 2023 19:50:39 +0300 Subject: [PATCH] IPEX Diffusers fix can't allocate 4GB+ with SDP --- modules/ipex_specific/diffusers.py | 162 ++++++++++++++++++++++++++++- modules/shared.py | 2 +- 2 files changed, 160 insertions(+), 4 deletions(-) diff --git a/modules/ipex_specific/diffusers.py b/modules/ipex_specific/diffusers.py index f8fab5a0f..f4314cb90 100644 --- a/modules/ipex_specific/diffusers.py +++ b/modules/ipex_specific/diffusers.py @@ -1,8 +1,10 @@ import torch import intel_extension_for_pytorch as ipex -import diffusers +import torch.nn.functional as F +import diffusers #1.19.3 + +Attention = diffusers.models.attention_processor.Attention -#ARC GPUs can't allocate more than 4GB to a single block: class SlicedAttnProcessor: r""" Processor for implementing sliced attention. @@ -16,7 +18,7 @@ class SlicedAttnProcessor: def __init__(self, slice_size): self.slice_size = slice_size - def __call__(self, attn: diffusers.models.attention_processor.Attention, hidden_states, encoder_hidden_states=None, attention_mask=None): + def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, attention_mask=None): residual = hidden_states input_ndim = hidden_states.ndim @@ -52,6 +54,7 @@ class SlicedAttnProcessor: (batch_size_attention, query_tokens, dim // attn.heads), device=query.device, dtype=query.dtype ) + #ARC GPUs can't allocate more than 4GB to a single block, Slice it: block_multiply = 2.4 if query.dtype == torch.float32 else 1.2 block_size = (batch_size_attention * query_tokens * shape_three) / 1024 * block_multiply #MB split_2_slice_size = query_tokens @@ -61,6 +64,9 @@ class SlicedAttnProcessor: sanity_check = 0 while ((self.slice_size * split_2_slice_size * shape_three) / 1024 * block_multiply) > 4000: split_2_slice_size = split_2_slice_size // 2 + if split_2_slice_size <= 1: + split_2_slice_size = 1 + break sanity_check = sanity_check + 1 if sanity_check >= 128: break @@ -112,5 +118,155 @@ class SlicedAttnProcessor: return hidden_states +class AttnProcessor2_0: + r""" + Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). + """ + + def __init__(self): + 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: Attention, + hidden_states, + encoder_hidden_states=None, + attention_mask=None, + temb=None, + ): + residual = hidden_states + + if attn.spatial_norm is not None: + hidden_states = attn.spatial_norm(hidden_states, temb) + + input_ndim = hidden_states.ndim + + if input_ndim == 4: + batch_size, channel, height, width = hidden_states.shape + hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) + + batch_size, sequence_length, _ = ( + hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + ) + + if attention_mask is not None: + attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + # scaled_dot_product_attention expects attention_mask shape to be + # (batch, heads, source_length, target_length) + attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) + + 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) + + 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) + value = attn.to_v(encoder_hidden_states) + + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + #ARC GPUs can't allocate more than 4GB to a single block, Slice it: + shape_one, batch_size_attention, query_tokens, shape_four = query.shape + block_multiply = 2.4 if query.dtype == torch.float32 else 1.2 + block_size = (shape_one * batch_size_attention * query_tokens * shape_four) / 1024 * block_multiply #MB + split_slice_size = batch_size_attention + if block_size >= 4000: + do_split = True + #Find something divisible with the shape_one + sanity_check = 0 + while ((shape_one * split_slice_size * query_tokens * shape_four) / 1024 * block_multiply) > 4000: + split_slice_size = split_slice_size // 2 + if split_slice_size <= 1: + split_slice_size = 1 + break + sanity_check = sanity_check + 1 + if sanity_check >= 128: + break + else: + do_split = False + + split_block_size = (shape_one * split_slice_size * query_tokens * shape_four) / 1024 * block_multiply #MB + split_2_slice_size = query_tokens + if split_block_size >= 4000: + do_split_2 = True + #Find something divisible with the batch_size_attention + sanity_check = 0 + while ((shape_one * split_slice_size * split_2_slice_size * shape_four) / 1024 * block_multiply) > 4000: + split_2_slice_size = split_2_slice_size // 2 + if split_2_slice_size <= 1: + split_2_slice_size = 1 + break + sanity_check = sanity_check + 1 + if sanity_check >= 128: + break + else: + do_split_2 = False + + if do_split: + hidden_states = torch.zeros(query.shape, device=query.device, dtype=query.dtype) + for i in range(batch_size_attention // split_slice_size): + start_idx = i * split_slice_size + end_idx = (i + 1) * split_slice_size + if do_split_2: + for i2 in range(query_tokens // split_2_slice_size): + start_idx_2 = i2 * split_2_slice_size + end_idx_2 = (i2 + 1) * split_2_slice_size + + query_slice = query[:, start_idx:end_idx, start_idx_2:end_idx_2] + key_slice = key[:, start_idx:end_idx, start_idx_2:end_idx_2] + attn_mask_slice = attention_mask[:, start_idx:end_idx, start_idx_2:end_idx_2] if attention_mask is not None else None + + attn_slice = F.scaled_dot_product_attention( + query_slice, key_slice, value[:, start_idx:end_idx, start_idx_2:end_idx_2], + attn_mask=attn_mask_slice, dropout_p=0.0, is_causal=False + ) + hidden_states[:, start_idx:end_idx, start_idx_2:end_idx_2] = attn_slice + else: + query_slice = query[:, start_idx:end_idx] + key_slice = key[:, start_idx:end_idx] + attn_mask_slice = attention_mask[:, start_idx:end_idx] if attention_mask is not None else None + + attn_slice = F.scaled_dot_product_attention( + query_slice, key_slice, value[:, start_idx:end_idx], + attn_mask=attn_mask_slice, dropout_p=0.0, is_causal=False + ) + hidden_states[:, start_idx:end_idx] = attn_slice + else: + hidden_states = F.scaled_dot_product_attention( + query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False + ) + + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + hidden_states = hidden_states.to(query.dtype) + + # linear proj + hidden_states = attn.to_out[0](hidden_states) + # dropout + hidden_states = attn.to_out[1](hidden_states) + + if input_ndim == 4: + hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) + + if attn.residual_connection: + hidden_states = hidden_states + residual + + hidden_states = hidden_states / attn.rescale_output_factor + + return hidden_states + def ipex_diffusers(): + #ARC GPUs can't allocate more than 4GB to a single block: diffusers.models.attention_processor.SlicedAttnProcessor = SlicedAttnProcessor + diffusers.models.attention_processor.AttnProcessor2_0 = AttnProcessor2_0 diff --git a/modules/shared.py b/modules/shared.py index 3dd194ff5..3c7952392 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -406,7 +406,7 @@ options_templates.update(options_section(('diffusers', "Diffusers Settings"), { "diffusers_vae_upcast": OptionInfo("default", "VAE upcasting", gr.Radio, lambda: {"choices": ['default', 'true', 'false']}), "diffusers_vae_slicing": OptionInfo(True, "Enable VAE slicing"), "diffusers_vae_tiling": OptionInfo(False, "Enable VAE tiling"), - "diffusers_attention_slicing": OptionInfo(True if devices.backend == "ipex" else False, "Enable attention slicing"), + "diffusers_attention_slicing": OptionInfo(False, "Enable attention slicing"), "diffusers_model_load_variant": OptionInfo("default", "Diffusers model loading variant", gr.Radio, lambda: {"choices": ['default', 'fp32', 'fp16']}), "diffusers_vae_load_variant": OptionInfo("default", "Diffusers VAE loading variant", gr.Radio, lambda: {"choices": ['default', 'fp32', 'fp16']}), # "diffusers_force_zeros": OptionInfo(False, "Force zeros for prompts when empty"),