From 8b40f475a31109cc6ecbdc0d14a0cee9e0303291 Mon Sep 17 00:00:00 2001
From: Nuullll <vfirst218@gmail.com>
Date: Fri, 10 Nov 2023 11:06:26 +0800
Subject: [PATCH] Initial IPEX support

---
 modules/devices.py      | 11 +++++++++--
 modules/xpu_specific.py | 42 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 2 deletions(-)
 create mode 100644 modules/xpu_specific.py

diff --git a/modules/devices.py b/modules/devices.py
index 1d4eb5635..be599736c 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -3,7 +3,7 @@ import contextlib
 from functools import lru_cache
 
 import torch
-from modules import errors, shared
+from modules import errors, shared, xpu_specific
 
 if sys.platform == "darwin":
     from modules import mac_specific
@@ -30,6 +30,9 @@ def get_optimal_device_name():
     if has_mps():
         return "mps"
 
+    if xpu_specific.has_ipex:
+        return xpu_specific.get_xpu_device_string()
+
     return "cpu"
 
 
@@ -100,11 +103,15 @@ def autocast(disable=False):
     if dtype == torch.float32 or shared.cmd_opts.precision == "full":
         return contextlib.nullcontext()
 
+    if xpu_specific.has_xpu:
+        return torch.autocast("xpu")
+
     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()
+    device_type = "xpu" if xpu_specific.has_xpu else "cuda"
+    return torch.autocast(device_type, enabled=False) if torch.is_autocast_enabled() and not disable else contextlib.nullcontext()
 
 
 class NansException(Exception):
diff --git a/modules/xpu_specific.py b/modules/xpu_specific.py
new file mode 100644
index 000000000..6417dd2d6
--- /dev/null
+++ b/modules/xpu_specific.py
@@ -0,0 +1,42 @@
+import contextlib
+from modules import shared
+from modules.sd_hijack_utils import CondFunc
+
+has_ipex = False
+try:
+    import torch
+    import intel_extension_for_pytorch as ipex
+    has_ipex = True
+except Exception:
+    pass
+
+def check_for_xpu():
+    if not has_ipex:
+        return False
+
+    return hasattr(torch, 'xpu') and torch.xpu.is_available()
+
+has_xpu = check_for_xpu()
+
+def get_xpu_device_string():
+    if shared.cmd_opts.device_id is not None:
+        return f"xpu:{shared.cmd_opts.device_id}"
+    return "xpu"
+
+def return_null_context(*args, **kwargs): # pylint: disable=unused-argument
+    return contextlib.nullcontext()
+
+if has_xpu:
+    CondFunc('torch.Generator',
+        lambda orig_func, device=None: torch.xpu.Generator(device),
+        lambda orig_func, device=None: device is not None and device != torch.device("cpu") and device != "cpu")
+
+    CondFunc('torch.nn.functional.layer_norm',
+        lambda orig_func, input, normalized_shape=None, weight=None, *args, **kwargs:
+        orig_func(input.to(weight.data.dtype), normalized_shape, weight, *args, **kwargs),
+        lambda orig_func, input, normalized_shape=None, weight=None, *args, **kwargs:
+        weight is not None and input.dtype != weight.data.dtype)
+
+    CondFunc('torch.nn.modules.GroupNorm.forward',
+        lambda orig_func, self, input: orig_func(self, input.to(self.weight.data.dtype)),
+        lambda orig_func, self, input: input.dtype != self.weight.data.dtype)