IPEX Diffusers fix cannot allocate more than 4GB

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':

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()

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

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:

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"),

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 "$@"