Intel ARC Support

modules/cmd_args.py

@@ -23,6 +23,7 @@ parser.add_argument("--allow-code", action='store_true', help="Allow custom scri
 parser.add_argument("--share", action='store_true', help="Enable to make the UI accessible through Gradio site")
 parser.add_argument("--enable-insecure", action='store_true', help="Enable extensions tab regardless of other options")
 parser.add_argument("--use-cpu", nargs='+', help="Force use CPU for specified modules", default=[], type=str.lower)
+parser.add_argument("--use-ipex", action='store_true', help="Force use Intel OneAPI XPU backend")
 parser.add_argument("--listen", action='store_true', help="Launch web server using public IP address")
 parser.add_argument("--port", type=int, help="Launch web server with given server port", default=None)
 parser.add_argument("--hide-ui-dir-config", action='store_true', help="Hide directory configuration from UI", default=False)

modules/codeformer_model.py

@@ -103,7 +103,11 @@ 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()
+                        from modules import shared
+                        if shared.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/devices.py

@@ -22,9 +22,13 @@ 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:
+        from modules import shared
+        if shared.cmd_opts.device_id is not None:
+            return f"cuda:{shared.cmd_opts.device_id}"
+        return "cuda"
 
 
 def get_dml_device_string():
@@ -35,7 +39,10 @@ def get_dml_device_string():
 
 
 def get_optimal_device_name():
-    if torch.cuda.is_available():
+    from modules import shared
+    if shared.cmd_opts.use_ipex:
+        return "xpu"
+    elif torch.cuda.is_available():
         return get_cuda_device_string()
     if has_mps():
         return "mps"
@@ -61,7 +68,11 @@ def get_device_for(task):
 
 
 def torch_gc():
-    if torch.cuda.is_available():
+    from modules import shared
+    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()
@@ -137,11 +148,19 @@ def autocast(disable=False):
         return contextlib.nullcontext()
     if dtype == torch.float32 or shared.cmd_opts.precision == "Full":
         return contextlib.nullcontext()
-    return torch.autocast("cuda")
+    from modules import shared
+    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()
+    from modules import shared
+    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):

modules/memmon.py

@@ -19,26 +19,44 @@ class MemUsageMonitor(threading.Thread):
         self.daemon = True
         self.run_flag = threading.Event()
         self.data = defaultdict(int)
-        if not torch.cuda.is_available():
+        from modules import shared
+        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)
+        from modules import shared
+        if shared.cmd_opts.use_ipex:
+            return torch.xpu.mem_get_info("xpu")
+        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()
+            from modules import shared
+            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 +72,19 @@ 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)
+        from modules import shared
+        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 +95,11 @@ class MemUsageMonitor(threading.Thread):
             self.data["free"] = free
             self.data["total"] = total
 
-            torch_stats = torch.cuda.memory_stats(self.device)
+            from modules import shared
+            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/processing.py

@@ -55,7 +55,17 @@ def memory_stats():
     except Exception as e:
         mem.update({ 'ram': e })
     try:
-        if torch.cuda.is_available():
+        from modules import shared
+        if shared.cmd_opts.use_ipex:
+            s = torch.xpu.mem_get_info()
+            gpu = { 'used': gb(s[1] - s[0]), 'total': gb(s[1]) }
+            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/sd_hijack_optimizations.py

@@ -22,7 +22,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.mem_get_info("xpu")
+        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 +197,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.mem_get_info("xpu")
+        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 +222,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 +418,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 +527,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

@@ -3,6 +3,7 @@ from packaging import version
 
 from modules import devices
 from modules.sd_hijack_utils import CondFunc
+from modules import shared
 
 
 class TorchHijackForUnet:
@@ -67,7 +68,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_models.py

@@ -533,7 +533,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/shared.py

@@ -238,7 +238,7 @@ options_templates.update(options_section(('sd', "Stable Diffusion"), {
     "comma_padding_backtrack": OptionInfo(20, "Increase coherency by padding from the last comma within n tokens when using more than 75 tokens", gr.Slider, {"minimum": 0, "maximum": 74, "step": 1 }),
     "CLIP_stop_at_last_layers": OptionInfo(1, "Clip skip", gr.Slider, {"minimum": 1, "maximum": 12, "step": 1, "visible": False}),
     "upcast_attn": OptionInfo(False, "Upcast cross attention layer to float32"),
-    "cross_attention_optimization": OptionInfo("Sub-quadratic" if is_device_dml else "Scaled-Dot-Product", "Cross-attention optimization method", gr.Radio, lambda: {"choices": shared_items.list_crossattention() }),
+    "cross_attention_optimization": OptionInfo("Sub-quadratic" if is_device_dml else "Split attention" if cmd_opts.use_ipex else "Scaled-Dot-Product", "Cross-attention optimization method", gr.Radio, lambda: {"choices": shared_items.list_crossattention() }),
     "cross_attention_options": OptionInfo([], "Cross-attention advanced options", gr.CheckboxGroup, lambda: {"choices": ['xFormers enable flash Attention', 'SDP disable memory attention']}),
     "sub_quad_q_chunk_size": OptionInfo(512, "Sub-quadratic cross-attention query chunk size for the  layer optimization to use", gr.Slider, {"minimum": 16, "maximum": 8192, "step": 8}),
     "sub_quad_kv_chunk_size": OptionInfo(512, "Sub-quadratic cross-attentionkv chunk size for the sub-quadratic cross-attention layer optimization to use", gr.Slider, {"minimum": 0, "maximum": 8192, "step": 8}),
@@ -318,7 +318,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/textual_inversion/textual_inversion.py

@@ -434,7 +434,11 @@ 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()
+    from modules import shared
+    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

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()}')
@@ -188,6 +194,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')
@@ -197,6 +204,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.nodirectml: # torch-directml is available on AMD64
@@ -212,7 +224,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: