Merge pull request #643 from Disty0/master

Intel ARC Support

modules/api/api.py

@@ -573,7 +573,24 @@ class Api:
             ram = { 'error': f'{err}' }
         try:
             import torch
-            if torch.cuda.is_available():
+            if shared.cmd_opts.use_ipex():
+                import intel_extension_for_pytorch as ipex
+                system = { 'free': (torch.xpu.get_device_properties("xpu").total_memory - torch.xpu.memory_allocated()), 'used': torch.xpu.memory_allocated(), 'total': torch.xpu.get_device_properties("xpu").total_memory }
+                s = dict(torch.xpu.memory_stats("xpu"))
+                allocated = { 'current': s['allocated_bytes.all.current'], 'peak': s['allocated_bytes.all.peak'] }
+                reserved = { 'current': s['reserved_bytes.all.current'], 'peak': s['reserved_bytes.all.peak'] }
+                active = { 'current': s['active_bytes.all.current'], 'peak': s['active_bytes.all.peak'] }
+                inactive = { 'current': s['inactive_split_bytes.all.current'], 'peak': s['inactive_split_bytes.all.peak'] }
+                warnings = { 'retries': s['num_alloc_retries'], 'oom': s['num_ooms'] }
+                cuda = {
+                    'system': system,
+                    'active': active,
+                    'allocated': allocated,
+                    'reserved': reserved,
+                    'inactive': inactive,
+                    'events': warnings,
+                }
+            elif torch.cuda.is_available():
                 s = torch.cuda.mem_get_info()
                 system = { 'free': s[0], 'used': s[1] - s[0], 'total': s[1] }
                 s = dict(torch.cuda.memory_stats(shared.device))

modules/cmd_args.py

@@ -27,6 +27,7 @@ group.add_argument("--allow-code", action='store_true', help="Allow custom scrip
 group.add_argument("--share", action='store_true', help="Enable UI accessible through Gradio site, default: %(default)s")
 group.add_argument("--insecure", action='store_true', help="Enable extensions tab regardless of other options, default: %(default)s")
 group.add_argument("--use-cpu", nargs='+', default=[], type=str.lower, help="Force use CPU for specified modules, default: %(default)s")
+group.add_argument("--use-ipex", action='store_true', help="Force use Intel OneAPI XPU backend, default: %(default)s", default=False)
 group.add_argument("--listen", action='store_true', help="Launch web server using public IP address, default: %(default)s")
 group.add_argument("--port", type=int, default=7860, help="Launch web server with given server port, default: %(default)s")
 group.add_argument("--freeze", action='store_true', help="Disable editing settings", default=False)

modules/codeformer/codeformer_arch.py

@@ -3,6 +3,10 @@
 import math
 import numpy as np
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 from torch import nn, Tensor
 import torch.nn.functional as F
 from typing import Optional, List

modules/codeformer/vqgan_arch.py

@@ -7,6 +7,10 @@ https://github.com/samb-t/unleashing-transformers/blob/master/models/vqgan.py
 '''
 import numpy as np
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 import torch.nn as nn
 import torch.nn.functional as F
 import copy

modules/codeformer_model.py

@@ -3,6 +3,10 @@ import sys
 
 import cv2
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 
 import modules.face_restoration
 from modules import shared, devices, modelloader, errors
@@ -103,7 +107,10 @@ def setup_model(dirname):
                             output = self.net(cropped_face_t, w=w if w is not None else shared.opts.code_former_weight, adain=True)[0]
                             restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
                         del output
-                        torch.cuda.empty_cache()
+                        if cmd_opts.use_ipex:
+                            torch.xpu.empty_cache()
+                        else:
+                            torch.cuda.empty_cache()
                     except Exception as error:
                         print(f'\tFailed inference for CodeFormer: {error}', file=sys.stderr)
                         restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))

modules/deepbooru.py

@@ -2,6 +2,10 @@ import os
 import re
 
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 import numpy as np
 
 from modules import modelloader, paths, deepbooru_model, devices, images, shared

modules/deepbooru_model.py

@@ -1,4 +1,8 @@
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 import torch.nn as nn
 import torch.nn.functional as F
 

modules/devices.py

@@ -1,6 +1,11 @@
 import sys
 import contextlib
 import torch
+from modules import shared
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 
 if sys.platform == "darwin":
     from modules import mac_specific
@@ -21,21 +26,24 @@ def extract_device_id(args, name):
 
 
 def get_cuda_device_string():
-    from modules import shared
-    if shared.cmd_opts.device_id is not None:
-        return f"cuda:{shared.cmd_opts.device_id}"
-    return "cuda"
+    if shared.cmd_opts.use_ipex:
+        return "xpu"
+    else:
+        if shared.cmd_opts.device_id is not None:
+            return f"cuda:{shared.cmd_opts.device_id}"
+        return "cuda"
 
 
 def get_dml_device_string():
-    from modules import shared
     if shared.cmd_opts.device_id is not None:
         return f"privateuseone:{shared.cmd_opts.device_id}"
     return "privateuseone:0"
 
 
 def get_optimal_device_name():
-    if torch.cuda.is_available():
+    if shared.cmd_opts.use_ipex:
+        return "xpu"
+    elif torch.cuda.is_available():
         return get_cuda_device_string()
     if has_mps():
         return "mps"
@@ -54,21 +62,22 @@ def get_optimal_device():
 
 
 def get_device_for(task):
-    from modules import shared
     if task in shared.cmd_opts.use_cpu:
         return cpu
     return get_optimal_device()
 
 
 def torch_gc():
-    if torch.cuda.is_available():
+    if shared.cmd_opts.use_ipex:
+        with torch.xpu.device("xpu"):
+            torch.xpu.empty_cache()
+    elif torch.cuda.is_available():
         with torch.cuda.device(get_cuda_device_string()):
             torch.cuda.empty_cache()
             torch.cuda.ipc_collect()
 
 
 def set_cuda_params():
-    from modules import shared
     if torch.cuda.is_available():
         try:
             torch.backends.cuda.matmul.allow_tf32 = shared.opts.cuda_allow_tf32
@@ -132,16 +141,21 @@ def randn_without_seed(shape):
 
 
 def autocast(disable=False):
-    from modules import shared
     if disable:
         return contextlib.nullcontext()
     if dtype == torch.float32 or shared.cmd_opts.precision == "Full":
         return contextlib.nullcontext()
-    return torch.autocast("cuda")
+    if shared.cmd_opts.use_ipex:
+        return torch.xpu.amp.autocast(enabled=True, dtype=dtype, cache_enabled=False)
+    else:
+        return torch.autocast("cuda")
 
 
 def without_autocast(disable=False):
-    return torch.autocast("cuda", enabled=False) if torch.is_autocast_enabled() and not disable else contextlib.nullcontext()
+    if shared.cmd_opts.use_ipex:
+        return torch.autocast("xpu", enabled=False) if torch.is_autocast_enabled() and not disable else contextlib.nullcontext()
+    else:
+        return torch.autocast("cuda", enabled=False) if torch.is_autocast_enabled() and not disable else contextlib.nullcontext()
 
 
 class NansException(Exception):
@@ -149,7 +163,6 @@ class NansException(Exception):
 
 
 def test_for_nans(x, where):
-    from modules import shared
     if shared.opts.disable_nan_check:
         return
     if not torch.all(torch.isnan(x)).item():

modules/esrgan_model.py

@@ -2,6 +2,10 @@ import os
 
 import numpy as np
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 from PIL import Image
 from basicsr.utils.download_util import load_file_from_url
 

modules/esrgan_model_arch.py

@@ -2,6 +2,10 @@
 
 import math
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 import torch.nn as nn
 import torch.nn.functional as F
 

modules/extras.py

@@ -4,6 +4,10 @@ import html
 import shutil
 
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 import tqdm
 import gradio as gr
 import safetensors.torch

modules/hypernetworks/hypernetwork.py

@@ -8,6 +8,10 @@ import inspect
 
 import modules.textual_inversion.dataset
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 import tqdm
 from einops import rearrange, repeat
 from ldm.util import default
@@ -591,7 +595,10 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi
             print("Cannot resume from saved optimizer!")
             print(e)
 
-    scaler = torch.cuda.amp.GradScaler()
+    if shared.cmd_opts.use_ipex:
+        scaler = torch.xpu.amp.GradScaler()
+    else:
+        scaler = torch.cuda.amp.GradScaler()
     
     batch_size = ds.batch_size
     gradient_step = ds.gradient_step
@@ -708,7 +715,9 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi
                     hypernetwork.eval()
                     rng_state = torch.get_rng_state()
                     cuda_rng_state = None
-                    if torch.cuda.is_available():
+                    if shared.cmd_opts.use_ipex:
+                        cuda_rng_state = torch.xpu.get_rng_state_all()
+                    elif torch.cuda.is_available():
                         cuda_rng_state = torch.cuda.get_rng_state_all()
                     shared.sd_model.cond_stage_model.to(devices.device)
                     shared.sd_model.first_stage_model.to(devices.device)
@@ -745,7 +754,9 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi
                         shared.sd_model.cond_stage_model.to(devices.cpu)
                         shared.sd_model.first_stage_model.to(devices.cpu)
                     torch.set_rng_state(rng_state)
-                    if torch.cuda.is_available():
+                    if shared.cmd_opts.use_ipex:
+                        torch.xpu.set_rng_state_all(cuda_rng_state)
+                    elif torch.cuda.is_available():
                         torch.cuda.set_rng_state_all(cuda_rng_state)
                     hypernetwork.train()
                     if image is not None:

modules/interrogate.py

@@ -5,6 +5,10 @@ from pathlib import Path
 import re
 
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 import torch.hub
 
 from torchvision import transforms

modules/lowvram.py

@@ -1,4 +1,8 @@
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 from modules import devices
 
 module_in_gpu = None

modules/mac_specific.py

@@ -1,4 +1,8 @@
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 import platform
 from modules.sd_hijack_utils import CondFunc
 from packaging import version

modules/memmon.py

@@ -2,6 +2,12 @@ import threading
 import time
 from collections import defaultdict
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
+
+from modules import shared
 
 
 class MemUsageMonitor(threading.Thread):
@@ -19,26 +25,41 @@ class MemUsageMonitor(threading.Thread):
         self.daemon = True
         self.run_flag = threading.Event()
         self.data = defaultdict(int)
-        if not torch.cuda.is_available():
+        if not torch.cuda.is_available() or not shared.cmd_opts.use_ipex:
             self.disabled = True
         else:
-            try:
-                self.cuda_mem_get_info()
-                torch.cuda.memory_stats(self.device)
-            except Exception as e:  # AMD or whatever
-                print(f"Torch exception: {e}")
-                self.disabled = True
+            if shared.cmd_opts.use_ipex:
+                try:
+                    self.cuda_mem_get_info()
+                    torch.cuda.memory_stats("xpu")
+                except Exception as e:  # AMD or whatever
+                    print(f"Torch exception: {e}")
+                    self.disabled = True
+
+            else:
+                try:
+                    self.cuda_mem_get_info()
+                    torch.cuda.memory_stats(self.device)
+                except Exception as e:  # AMD or whatever
+                    print(f"Torch exception: {e}")
+                    self.disabled = True
 
     def cuda_mem_get_info(self):
-        index = self.device.index if self.device.index is not None else torch.cuda.current_device()
-        return torch.cuda.mem_get_info(index)
+        if shared.cmd_opts.use_ipex:
+            return [(torch.xpu.get_device_properties("xpu").total_memory - torch.xpu.memory_allocated()), torch.xpu.get_device_properties("xpu").total_memory]
+        else:
+            index = self.device.index if self.device.index is not None else torch.cuda.current_device()
+            return torch.cuda.mem_get_info(index)
 
     def run(self):
         if self.disabled:
             return
         while True:
             self.run_flag.wait()
-            torch.cuda.reset_peak_memory_stats()
+            if shared.cmd_opts.use_ipex:
+                torch.xpu.reset_peak_memory_stats()
+            else:
+                torch.cuda.reset_peak_memory_stats()
             self.data.clear()
             if self.opts.memmon_poll_rate <= 0:
                 self.run_flag.clear()
@@ -54,12 +75,18 @@ class MemUsageMonitor(threading.Thread):
         for k, v in self.read().items():
             print(k, -(v // -(1024 ** 2)))
         print(self, 'raw torch memory stats:')
-        tm = torch.cuda.memory_stats(self.device)
+        if shared.cmd_opts.use_ipex:
+            tm = torch.xpu.memory_stats("xpu")
+        else:
+            tm = torch.cuda.memory_stats(self.device)
         for k, v in tm.items():
             if 'bytes' not in k:
                 continue
             print('\t' if 'peak' in k else '', k, -(v // -(1024 ** 2)))
-        print(torch.cuda.memory_summary())
+        if shared.cmd_opts.use_ipex:
+            print(torch.xpu.memory_summary())
+        else:
+            print(torch.cuda.memory_summary())
 
     def monitor(self):
         self.run_flag.set()
@@ -70,7 +97,10 @@ class MemUsageMonitor(threading.Thread):
             self.data["free"] = free
             self.data["total"] = total
 
-            torch_stats = torch.cuda.memory_stats(self.device)
+            if shared.cmd_opts.use_ipex:
+                torch_stats = torch.xpu.memory_stats("xpu")
+            else:
+                torch_stats = torch.cuda.memory_stats(self.device)
             self.data["active"] = torch_stats["active.all.current"]
             self.data["active_peak"] = torch_stats["active_bytes.all.peak"]
             self.data["reserved"] = torch_stats["reserved_bytes.all.current"]

modules/models/diffusion/ddpm_edit.py

@@ -10,6 +10,10 @@ https://github.com/CompVis/taming-transformers
 # See more details in LICENSE.
 
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 import torch.nn as nn
 import numpy as np
 import pytorch_lightning as pl

modules/models/diffusion/uni_pc/sampler.py

@@ -2,6 +2,10 @@
 
 import numpy as np
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 
 from .uni_pc import NoiseScheduleVP, model_wrapper, UniPC
 from modules import shared, devices

modules/models/diffusion/uni_pc/uni_pc.py

@@ -1,4 +1,8 @@
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 import torch.nn.functional as F
 import math
 import time

modules/processing.py

@@ -8,6 +8,10 @@ from typing import Any, Dict, List
 
 import psutil
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 import numpy as np
 from PIL import Image, ImageFilter, ImageOps
 import cv2
@@ -55,7 +59,15 @@ def memory_stats():
     except Exception as e:
         mem.update({ 'ram': e })
     try:
-        if torch.cuda.is_available():
+        if cmd_opts.use_ipex:
+            gpu = { 'used': gb(torch.xpu.memory_allocated()), 'total': gb(torch.xpu.get_device_properties("xpu").total_memory) }
+            s = dict(torch.xpu.memory_stats("xpu"))
+            mem.update({
+                'gpu': gpu,
+                'retries': s['num_alloc_retries'],
+                'oom': s['num_ooms']
+            })
+        elif torch.cuda.is_available():
             s = torch.cuda.mem_get_info()
             gpu = { 'used': gb(s[1] - s[0]), 'total': gb(s[1]) }
             s = dict(torch.cuda.memory_stats(shared.device))

modules/prompt_parser.py

@@ -368,3 +368,7 @@ if __name__ == "__main__":
     doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)
 else:
     import torch  # doctest faster
+    try:
+        import intel_extension_for_pytorch as ipex
+    except:
+        pass

modules/safe.py

@@ -6,6 +6,10 @@ import zipfile
 import re
 
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 import numpy
 import _codecs
 

modules/sd_disable_initialization.py

@@ -1,6 +1,10 @@
 import ldm.modules.encoders.modules
 import open_clip
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 import transformers.utils.hub
 
 

modules/sd_hijack.py

@@ -1,6 +1,10 @@
 from types import MethodType
 from rich import print # pylint: disable=redefined-builtin
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 from torch.nn.functional import silu
 import ldm.modules.attention
 import ldm.modules.diffusionmodules.model

modules/sd_hijack_clip.py

@@ -2,6 +2,10 @@ import math
 from collections import namedtuple
 
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 
 from modules import prompt_parser, devices, sd_hijack
 from modules.shared import opts

modules/sd_hijack_inpainting.py

@@ -1,4 +1,8 @@
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 
 import ldm.models.diffusion.ddpm
 import ldm.models.diffusion.ddim

modules/sd_hijack_open_clip.py

@@ -1,5 +1,9 @@
 import open_clip.tokenizer
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 
 from modules import sd_hijack_clip, devices
 

modules/sd_hijack_optimizations.py

@@ -2,6 +2,10 @@ import math
 import psutil
 
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 from torch import einsum
 
 from ldm.util import default
@@ -22,7 +26,15 @@ if shared.opts.cross_attention_optimization == "xFormers":
 
 
 def get_available_vram():
-    if shared.device.type == 'cuda':
+    if shared.cmd_opts.use_ipex:
+        stats = torch.xpu.memory_stats("xpu")
+        mem_active = stats['active_bytes.all.current']
+        mem_reserved = stats['reserved_bytes.all.current']
+        mem_free_xpu, _ = [(torch.xpu.get_device_properties("xpu").total_memory - torch.xpu.memory_allocated()), torch.xpu.get_device_properties("xpu").total_memory]
+        mem_free_torch = mem_reserved - mem_active
+        mem_free_total = mem_free_xpu + mem_free_torch
+        return mem_free_total
+    elif shared.device.type == 'cuda':
         stats = torch.cuda.memory_stats(shared.device)
         mem_active = stats['active_bytes.all.current']
         mem_reserved = stats['reserved_bytes.all.current']
@@ -189,14 +201,24 @@ def einsum_op_tensor_mem(q, k, v, max_tensor_mb):
     return einsum_op_slice_1(q, k, v, max(q.shape[1] // div, 1))
 
 def einsum_op_cuda(q, k, v):
-    stats = torch.cuda.memory_stats(q.device)
-    mem_active = stats['active_bytes.all.current']
-    mem_reserved = stats['reserved_bytes.all.current']
-    mem_free_cuda, _ = torch.cuda.mem_get_info(q.device)
-    mem_free_torch = mem_reserved - mem_active
-    mem_free_total = mem_free_cuda + mem_free_torch
-    # Divide factor of safety as there's copying and fragmentation
-    return einsum_op_tensor_mem(q, k, v, mem_free_total / 3.3 / (1 << 20))
+    if shared.cmd_opts.use_ipex:
+        stats = torch.xpu.memory_stats("xpu")
+        mem_active = stats['active_bytes.all.current']
+        mem_reserved = stats['reserved_bytes.all.current']
+        mem_free_xpu, _ = [(torch.xpu.get_device_properties("xpu").total_memory - torch.xpu.memory_allocated()), torch.xpu.get_device_properties("xpu").total_memory]
+        mem_free_torch = mem_reserved - mem_active
+        mem_free_total = mem_free_xpu + mem_free_torch
+        # Divide factor of safety as there's copying and fragmentation
+        return einsum_op_tensor_mem(q, k, v, mem_free_total / 3.3 / (1 << 20))
+    else:
+        stats = torch.cuda.memory_stats(q.device)
+        mem_active = stats['active_bytes.all.current']
+        mem_reserved = stats['reserved_bytes.all.current']
+        mem_free_cuda, _ = torch.cuda.mem_get_info(q.device)
+        mem_free_torch = mem_reserved - mem_active
+        mem_free_total = mem_free_cuda + mem_free_torch
+        # Divide factor of safety as there's copying and fragmentation
+        return einsum_op_tensor_mem(q, k, v, mem_free_total / 3.3 / (1 << 20))
 
 def einsum_op_dml(q, k, v):
     mem_total, mem_active = torch.dml.memory_stats(q.device)
@@ -204,6 +226,9 @@ def einsum_op_dml(q, k, v):
     return einsum_op_tensor_mem(q, k, v, (mem_reserved - mem_active) if mem_reserved > mem_active else 1)
 
 def einsum_op(q, k, v):
+    if shared.cmd_opts.use_ipex:
+        return einsum_op_cuda(q, k, v)
+
     if q.device.type == 'cuda':
         return einsum_op_cuda(q, k, v)
 
@@ -397,8 +422,12 @@ def scaled_dot_product_attention_forward(self, x, context=None, mask=None):
     return hidden_states
 
 def scaled_dot_product_no_mem_attention_forward(self, x, context=None, mask=None):
-    with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False):
-        return scaled_dot_product_attention_forward(self, x, context, mask)
+    if shared.cmd_opts.use_ipex:
+        with torch.backends.xpu.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False):
+            return scaled_dot_product_attention_forward(self, x, context, mask)
+    else:
+        with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False):
+            return scaled_dot_product_attention_forward(self, x, context, mask)
 
 def cross_attention_attnblock_forward(self, x):
     h_ = x
@@ -502,8 +531,12 @@ def sdp_attnblock_forward(self, x):
     return x + out
 
 def sdp_no_mem_attnblock_forward(self, x):
-    with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False):
-        return sdp_attnblock_forward(self, x)
+    if shared.cmd_opts.use_ipex:
+        with torch.backends.xpu.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False):
+            return sdp_attnblock_forward(self, x)
+    else:
+        with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False):
+            return sdp_attnblock_forward(self, x)
 
 def sub_quad_attnblock_forward(self, x):
     h_ = x

modules/sd_hijack_unet.py

@@ -1,8 +1,13 @@
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 from packaging import version
 
 from modules import devices
 from modules.sd_hijack_utils import CondFunc
+from modules import shared
 
 
 class TorchHijackForUnet:
@@ -67,7 +72,7 @@ def hijack_ddpm_edit():
 unet_needs_upcast = lambda *args, **kwargs: devices.unet_needs_upcast
 CondFunc('ldm.models.diffusion.ddpm.LatentDiffusion.apply_model', apply_model, unet_needs_upcast)
 CondFunc('ldm.modules.diffusionmodules.openaimodel.timestep_embedding', lambda orig_func, timesteps, *args, **kwargs: orig_func(timesteps, *args, **kwargs).to(torch.float32 if timesteps.dtype == torch.int64 else devices.dtype_unet), unet_needs_upcast)
-if version.parse(torch.__version__) <= version.parse("1.13.2") or torch.cuda.is_available():
+if version.parse(torch.__version__) <= version.parse("1.13.2") or torch.cuda.is_available() or shared.cmd_opts.use_ipex:
     CondFunc('ldm.modules.diffusionmodules.util.GroupNorm32.forward', lambda orig_func, self, *args, **kwargs: orig_func(self.float(), *args, **kwargs), unet_needs_upcast)
     CondFunc('ldm.modules.attention.GEGLU.forward', lambda orig_func, self, x: orig_func(self.float(), x.float()).to(devices.dtype_unet), unet_needs_upcast)
     CondFunc('open_clip.transformer.ResidualAttentionBlock.__init__', lambda orig_func, *args, **kwargs: kwargs.update({'act_layer': GELUHijack}) and False or orig_func(*args, **kwargs), lambda _, *args, **kwargs: kwargs.get('act_layer') is None or kwargs['act_layer'] == torch.nn.GELU)

modules/sd_hijack_xlmr.py

@@ -1,4 +1,8 @@
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 
 from modules import sd_hijack_clip, devices
 

modules/sd_models.py

@@ -8,6 +8,10 @@ from os import mkdir
 from urllib import request
 from rich import print, progress # pylint: disable=redefined-builtin
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 import safetensors.torch
 from omegaconf import OmegaConf
 import tomesd
@@ -533,7 +537,6 @@ def unload_model_weights(sd_model=None, _info=None):
         sd_model = None
         gc.collect()
         devices.torch_gc()
-        torch.cuda.empty_cache()
     print(f"Unloaded weights {timer.summary()}")
     return sd_model
 

modules/sd_models_config.py

@@ -1,6 +1,10 @@
 import os
 
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 
 from modules import paths, sd_disable_initialization
 

modules/sd_samplers_common.py

@@ -1,6 +1,10 @@
 from collections import namedtuple
 import numpy as np
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 from PIL import Image
 from modules import devices, processing, images, sd_vae_approx
 

modules/sd_samplers_compvis.py

@@ -4,6 +4,10 @@ import ldm.models.diffusion.plms
 
 import numpy as np
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 
 from modules.shared import state
 from modules import sd_samplers_common, prompt_parser, shared

modules/sd_samplers_kdiffusion.py

@@ -1,6 +1,10 @@
 from collections import deque
 import inspect
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 import k_diffusion.sampling
 from modules import prompt_parser, devices, sd_samplers_common
 

modules/sd_vae.py

@@ -3,8 +3,14 @@ import collections
 import glob
 from copy import deepcopy
 from rich import print # pylint: disable=redefined-builtin
+from modules import shared
 import torch
-from modules import paths, shared, devices, script_callbacks, sd_models
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    if shared.cmd_opts.use_ipex:
+        print("Failed to import IPEX")
+from modules import paths, devices, script_callbacks, sd_models
 
 
 vae_ignore_keys = {"model_ema.decay", "model_ema.num_updates"}

modules/sd_vae_approx.py

@@ -1,6 +1,10 @@
 import os
 
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 from torch import nn
 from modules import devices, paths
 

modules/shared.py

@@ -316,7 +316,7 @@ options_templates.update(options_section(('cuda', "CUDA Settings"), {
     "cuda_dtype": OptionInfo("FP32" if sys.platform == "darwin" else "FP16", "Device precision type", gr.Radio, lambda: {"choices": ["FP32", "FP16", "BF16"]}),
     "no_half": OptionInfo(True if is_device_dml else False, "Use full precision for model (--no-half)", None, None, lambda: print("Warning: Most of DirectML devices do not fully support half mode. Recommend to use full precision to model.") if is_device_dml else None),
     "no_half_vae": OptionInfo(True if is_device_dml else False, "Use full precision for VAE (--no-half-vae)"),
-    "upcast_sampling": OptionInfo(True if sys.platform == "darwin" else False, "Enable upcast sampling. Usually produces similar results to --no-half with better performance while using less memory"),
+    "upcast_sampling": OptionInfo(True if sys.platform == "darwin" or cmd_opts.use_ipex else False, "Enable upcast sampling. Usually produces similar results to --no-half with better performance while using less memory"),
     "disable_nan_check": OptionInfo(True, "Do not check if produced images/latent spaces have NaN values"),
     "rollback_vae": OptionInfo(False, "Attempt to roll back VAE when produced NaN values, requires NaN check (experimental)"),
     "opt_channelslast": OptionInfo(False, "Use channels last as torch memory format "),

modules/sub_quadratic_attention.py

@@ -14,6 +14,10 @@ from functools import partial
 import math
 from typing import Optional, NamedTuple, List
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 from torch import Tensor
 from torch.utils.checkpoint import checkpoint
 

modules/textual_inversion/dataset.py

@@ -2,6 +2,10 @@ import os
 import numpy as np
 import PIL
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 from PIL import Image
 from torch.utils.data import Dataset, DataLoader, Sampler
 from torchvision import transforms

modules/textual_inversion/image_embedding.py

@@ -4,6 +4,10 @@ import numpy as np
 import zlib
 from PIL import Image, PngImagePlugin, ImageDraw, ImageFont
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 from modules.shared import opts
 
 

modules/textual_inversion/textual_inversion.py

@@ -3,6 +3,10 @@ import html
 import csv
 from collections import namedtuple
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 import tqdm
 import safetensors.torch
 from rich import print # pylint: disable=redefined-builtin
@@ -435,7 +439,10 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st
         else:
             print("No saved optimizer exists in checkpoint")
 
-    scaler = torch.cuda.amp.GradScaler()
+    if shared.cmd_opts.use_ipex:
+        scaler = torch.xpu.amp.GradScaler()
+    else:
+        scaler = torch.cuda.amp.GradScaler()
 
     batch_size = ds.batch_size
     gradient_step = ds.gradient_step

modules/xlmr.py

@@ -1,5 +1,9 @@
 from typing import Optional
 import torch
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 import torch.nn as nn
 from transformers import XLMRobertaModel,XLMRobertaTokenizer, BertPreTrainedModel, BertModel, BertConfig # pylint: disable=unused-import
 from transformers.models.xlm_roberta.configuration_xlm_roberta import XLMRobertaConfig

setup.py

@@ -56,6 +56,7 @@ def setup_logging(clean=False):
 # check if package is installed
 def installed(package, friendly: str = None):
     import pkg_resources
+    from modules import shared
     ok = True
     try:
         if friendly:
@@ -76,6 +77,8 @@ def installed(package, friendly: str = None):
             ok = ok and spec is not None
             if ok:
                 version = pkg_resources.get_distribution(p[0]).version
+                if shared.cmd_opts.use_ipex and p[0] == "pytorch_lightning":
+                    p[1] = "1.8.6"
                 log.debug(f"Package version found: {p[0]} {version}")
                 if len(p) > 1:
                     ok = ok and version == p[1]
@@ -91,6 +94,9 @@ def installed(package, friendly: str = None):
 
 # install package using pip if not already installed
 def install(package, friendly: str = None, ignore: bool = False):
+    from modules import shared
+    if shared.cmd_opts.use_ipex and package == "pytorch_lightning==1.9.4":
+        package = "pytorch_lightning==1.8.6"
     def pip(arg: str):
         arg = arg.replace('>=', '==')
         log.info(f'Installing package: {arg.replace("install", "").replace("--upgrade", "").replace("--no-deps", "").replace("  ", " ").strip()}')
@@ -187,6 +193,7 @@ def check_python():
 
 # check torch version
 def check_torch():
+    from modules import shared
     if shutil.which('nvidia-smi') is not None or os.path.exists(os.path.join(os.environ.get('SystemRoot') or r'C:\Windows', 'System32', 'nvidia-smi.exe')):
         log.info('nVidia toolkit detected')
         torch_command = os.environ.get('TORCH_COMMAND', 'torch torchaudio torchvision --index-url https://download.pytorch.org/whl/cu118')
@@ -196,6 +203,11 @@ def check_torch():
         os.environ.setdefault('HSA_OVERRIDE_GFX_VERSION', '10.3.0')
         torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm5.4.2')
         xformers_package = os.environ.get('XFORMERS_PACKAGE', 'none')
+    elif shutil.which('sycl-ls') is not None or os.path.exists('/opt/intel/oneapi'):
+        shared.cmd_opts.use_ipex = True
+        log.info('Intel toolkit detected')
+        torch_command = os.environ.get('TORCH_COMMAND', 'torch==1.13.0a0+git6c9b55e torchvision==0.14.1a0 intel_extension_for_pytorch==1.13.120+xpu --index-url https://developer.intel.com/ipex-whl-stable-xpu')
+        xformers_package = os.environ.get('XFORMERS_PACKAGE', 'none')
     else:
         machine = platform.machine()
         if 'arm' not in machine and 'aarch' not in machine and not args.no_directml: # torch-directml is available on AMD64
@@ -211,7 +223,11 @@ def check_torch():
     try:
         import torch
         log.info(f'Torch {torch.__version__}')
-        if torch.cuda.is_available():
+        if shared.cmd_opts.use_ipex:
+                import intel_extension_for_pytorch as ipex
+                log.info(f'Torch backend: Intel OneAPI {torch.__version__}')
+                log.info(f'Torch detected GPU: {torch.xpu.get_device_name("xpu")} VRAM {round(torch.xpu.get_device_properties("xpu").total_memory / 1024 / 1024)}')
+        elif torch.cuda.is_available():
             if torch.version.cuda:
                 log.info(f'Torch backend: nVidia CUDA {torch.version.cuda} cuDNN {torch.backends.cudnn.version() if torch.backends.cudnn.is_available() else "N/A"}')
             elif torch.version.hip:

webui.py

@@ -12,6 +12,10 @@ from modules import timer, errors
 startup_timer = timer.Timer()
 
 import torch # pylint: disable=C0411
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
 import torchvision # pylint: disable=W0611,C0411
 import pytorch_lightning # pytorch_lightning should be imported after torch, but it re-enables warnings on import so import once to disable them # pylint: disable=W0611,C0411
 logging.getLogger("xformers").addFilter(lambda record: 'A matching Triton is not available' not in record.getMessage())