From 489d0382cfd95ba4420d94dce4d2432c0cb54dd8 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sat, 5 Aug 2023 17:26:18 +0300 Subject: [PATCH] IPEX Diffusers fix cannot allocate more than 4GB --- installer.py | 2 +- modules/ipex_specific/__init__.py | 3 + modules/ipex_specific/diffusers.py | 111 +++++++++++++++++++++++++++++ modules/sd_vae.py | 7 +- modules/shared.py | 2 +- webui.sh | 2 +- 6 files changed, 121 insertions(+), 6 deletions(-) create mode 100644 modules/ipex_specific/diffusers.py diff --git a/installer.py b/installer.py index b1b5d1520..644b6b4b8 100644 --- a/installer.py +++ b/installer.py @@ -328,7 +328,7 @@ def check_torch(): torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.0.1a0 torchvision==0.15.2a0 intel_extension_for_pytorch==2.0.110+xpu -f https://developer.intel.com/ipex-whl-stable-xpu') os.environ.setdefault('TENSORFLOW_PACKAGE', 'tensorflow==2.13.0 intel-extension-for-tensorflow[gpu]') else: - torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.0.0a0 torchvision intel_extension_for_pytorch==2.0.110+gitba7f6c1 -f https://developer.intel.com/ipex-whl-stable-xpu') + torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.0.0a0 intel_extension_for_pytorch==2.0.110+gitba7f6c1 -f https://developer.intel.com/ipex-whl-stable-xpu') else: machine = platform.machine() if sys.platform == 'darwin': diff --git a/modules/ipex_specific/__init__.py b/modules/ipex_specific/__init__.py index b71027b35..8dbec2830 100644 --- a/modules/ipex_specific/__init__.py +++ b/modules/ipex_specific/__init__.py @@ -4,6 +4,7 @@ import torch import intel_extension_for_pytorch as ipex from modules import shared from modules.sd_hijack_utils import CondFunc +from .diffusers import ipex_diffusers #ControlNet depth_leres++ class DummyDataParallel(torch.nn.Module): @@ -149,3 +150,5 @@ def ipex_init(): weight if weight is not None else torch.ones(input.size()[1], device=shared.device), bias if bias is not None else torch.zeros(input.size()[1], device=shared.device), *args, **kwargs), lambda orig_func, input, *args, **kwargs: input.device != torch.device("cpu")) + + ipex_diffusers() diff --git a/modules/ipex_specific/diffusers.py b/modules/ipex_specific/diffusers.py new file mode 100644 index 000000000..7253e8cd5 --- /dev/null +++ b/modules/ipex_specific/diffusers.py @@ -0,0 +1,111 @@ +import torch +import intel_extension_for_pytorch as ipex +import diffusers + +#ARC GPUs can't allocate more than 4GB to a single block: +class SlicedAttnProcessor: + r""" + Processor for implementing sliced attention. + + Args: + slice_size (`int`, *optional*): + The number of steps to compute attention. Uses as many slices as `attention_head_dim // slice_size`, and + `attention_head_dim` must be a multiple of the `slice_size`. + """ + + 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): + residual = hidden_states + + 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 + ) + attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + + 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) + dim = query.shape[-1] + query = attn.head_to_batch_dim(query) + + 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) + key = attn.head_to_batch_dim(key) + value = attn.head_to_batch_dim(value) + + batch_size_attention, query_tokens, shape_three = query.shape + hidden_states = torch.zeros( + (batch_size_attention, query_tokens, dim // attn.heads), device=query.device, dtype=query.dtype + ) + + block_size = (batch_size_attention * query_tokens * shape_three) / 1024 * 1.2 #MB + split_2_slice_size = query_tokens + if block_size >= 4000: + do_split_2 = True + #Find something divisible with the query_tokens + while ((self.slice_size * split_2_slice_size * shape_three) / 1024 * 1.2) > 4000: + split_2_slice_size = split_2_slice_size // 2 + else: + do_split_2 = False + + for i in range(batch_size_attention // self.slice_size): + start_idx = i * self.slice_size + end_idx = (i + 1) * self.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 = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice) + attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx, start_idx_2:end_idx_2]) + + 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 = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice) + + attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx]) + + hidden_states[start_idx:end_idx] = attn_slice + + hidden_states = attn.batch_to_head_dim(hidden_states) + + # 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(): + diffusers.models.attention_processor.SlicedAttnProcessor = SlicedAttnProcessor diff --git a/modules/sd_vae.py b/modules/sd_vae.py index e87af749f..1edd39961 100644 --- a/modules/sd_vae.py +++ b/modules/sd_vae.py @@ -196,9 +196,10 @@ def load_vae_diffusers(_model, vae_file=None, vae_source="from unknown source"): try: import diffusers if os.path.isfile(vae_file): - # load_config passed to from_single_file doesn't apply - # from_single_file by default downloads VAE1.5 config - shared.log.warning("Using SDXL VAE loaded from singular file will result in low contrast images.") + if shared.opts.diffusers_pipeline == "Stable Diffusion XL": + # load_config passed to from_single_file doesn't apply + # from_single_file by default downloads VAE1.5 config + shared.log.warning("Using SDXL VAE loaded from singular file will result in low contrast images.") vae = diffusers.AutoencoderKL.from_single_file(vae_file) vae = vae.to(devices.dtype_vae) else: diff --git a/modules/shared.py b/modules/shared.py index 8b0ff3109..eb8c685c0 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -407,7 +407,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(False, "Enable attention slicing"), + "diffusers_attention_slicing": OptionInfo(True if devices.backend == "ipex" else 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"), diff --git a/webui.sh b/webui.sh index 550008f2d..2800b2d3d 100755 --- a/webui.sh +++ b/webui.sh @@ -96,7 +96,7 @@ if [[ ! -z "${ACCELERATE}" ]] && [ ${ACCELERATE}="True" ] && [ -x "$(command -v then echo "Launching accelerate launch.py..." exec accelerate launch --num_cpu_threads_per_process=6 launch.py "$@" -elif [[ -z "${first_launch}" ]] && [[ $(uname -a) != *WSL2* ]] && [ -x "$(command -v ipexrun)" ] && [ -x "$(command -v numactl)" ] && [ -x "$(command -v sycl-ls)" ] +elif [[ "$@" == *"--use-ipex"* ]] && [[ -z "${first_launch}" ]] && [[ $(uname -a) != *WSL2* ]] && [ -x "$(command -v ipexrun)" ] && [ -x "$(command -v numactl)" ] && [ -x "$(command -v sycl-ls)" ] then echo "Launching ipexrun launch.py..." exec ipexrun launch.py "$@"