From 065155f5e632232ca78d02b73fb1f29100282d8d Mon Sep 17 00:00:00 2001
From: Disty0 <disty@disty.xyz>
Date: Fri, 18 Oct 2024 01:12:28 +0300
Subject: [PATCH] IPEX fixes

---
 modules/apg/__init__.py        | 5 ++++-
 modules/intel/ipex/__init__.py | 2 +-
 modules/intel/ipex/hijacks.py  | 5 +++--
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/modules/apg/__init__.py b/modules/apg/__init__.py
index 0bf6f1c8d..1dc7a619e 100644
--- a/modules/apg/__init__.py
+++ b/modules/apg/__init__.py
@@ -27,12 +27,15 @@ def project(
     v0: torch.Tensor, # [B, C, H, W]
     v1: torch.Tensor, # [B, C, H, W]
     ):
+    device = v0.device
     dtype = v0.dtype
+    if device.type == "xpu":
+        v0, v1 = v0.to("cpu"), v1.to("cpu")
     v0, v1 = v0.double(), v1.double()
     v1 = torch.nn.functional.normalize(v1, dim=[-1, -2, -3])
     v0_parallel = (v0 * v1).sum(dim=[-1, -2, -3], keepdim=True) * v1
     v0_orthogonal = v0 - v0_parallel
-    return v0_parallel.to(dtype), v0_orthogonal.to(dtype)
+    return v0_parallel.to(device, dtype=dtype), v0_orthogonal.to(device, dtype=dtype)
 
 
 def normalized_guidance(
diff --git a/modules/intel/ipex/__init__.py b/modules/intel/ipex/__init__.py
index 49ca2c46f..bdc98d806 100644
--- a/modules/intel/ipex/__init__.py
+++ b/modules/intel/ipex/__init__.py
@@ -199,7 +199,7 @@ def ipex_init(): # pylint: disable=too-many-statements
             torch.cuda.ipc_collect = lambda *args, **kwargs: None
             torch.cuda.utilization = lambda *args, **kwargs: 0
 
-            ipex_hijacks()
+            ipex_hijacks(legacy=legacy)
             try:
                 from .diffusers import ipex_diffusers
                 ipex_diffusers()
diff --git a/modules/intel/ipex/hijacks.py b/modules/intel/ipex/hijacks.py
index bb79053c4..7ec94138d 100644
--- a/modules/intel/ipex/hijacks.py
+++ b/modules/intel/ipex/hijacks.py
@@ -293,7 +293,9 @@ def torch_load(f, map_location=None, *args, **kwargs):
 
 
 # Hijack Functions:
-def ipex_hijacks():
+def ipex_hijacks(legacy=True):
+    if legacy:
+        torch.nn.functional.interpolate = interpolate
     torch.tensor = torch_tensor
     torch.Tensor.to = Tensor_to
     torch.Tensor.cuda = Tensor_cuda
@@ -319,7 +321,6 @@ def ipex_hijacks():
     torch.nn.functional.layer_norm = functional_layer_norm
     torch.nn.functional.linear = functional_linear
     torch.nn.functional.conv2d = functional_conv2d
-    torch.nn.functional.interpolate = interpolate
     torch.nn.functional.pad = functional_pad
 
     torch.bmm = torch_bmm