This commit is contained in:
Disty0
2025-09-17 19:58:21 +03:00
parent cc51f79508
commit 2842f5e3aa
+24 -8
View File
@@ -1,5 +1,3 @@
from typing import Optional
import os
from functools import wraps
from contextlib import nullcontext
@@ -18,28 +16,27 @@ device_supports_fp64 = torch.xpu.has_fp64_dtype() if hasattr(torch.xpu, "has_fp6
# pylint: disable=protected-access, missing-function-docstring, line-too-long, unnecessary-lambda, no-else-return
class DummyDataParallel(torch.nn.Module): # pylint: disable=missing-class-docstring, unused-argument, too-few-public-methods
def __new__(cls, module, device_ids=None, output_device=None, dim=0): # pylint: disable=unused-argument
if isinstance(device_ids, list) and len(device_ids) > 1:
errors.log.error("IPEX backend doesn't support DataParallel on multiple XPU devices")
return module.to(devices.device)
def return_null_context(*args, **kwargs): # pylint: disable=unused-argument
return nullcontext()
@property
def is_cuda(self):
return self.device.type == "xpu" or self.device.type == "cuda"
def check_device_type(device, device_type: str) -> bool:
if device is None or type(device) not in {str, int, torch.device}:
return False
else:
return bool(torch.device(device).type == device_type)
def check_cuda(device) -> bool:
return bool(isinstance(device, int) or check_device_type(device, "cuda"))
def return_xpu(device): # keep the device instance type, aka return string if the input is string
return devices.device if device is None else f"xpu:{device.split(':')[-1]}" if isinstance(device, str) and ":" in device else f"xpu:{device}" if isinstance(device, int) else torch.device(f"xpu:{device.index}" if device.index is not None else "xpu") if isinstance(device, torch.device) else "xpu"
@@ -107,6 +104,7 @@ def functional_pad(input, pad, mode='constant', value=None):
else:
return original_functional_pad(input, pad, mode=mode, value=value)
# Diffusers FreeU
original_fft_fftn = torch.fft.fftn
@wraps(torch.fft.fftn)
@@ -114,6 +112,7 @@ def fft_fftn(input, s=None, dim=None, norm=None, *, out=None):
return_dtype = input.dtype
return original_fft_fftn(input.to(dtype=torch.float32), s=s, dim=dim, norm=norm, out=out).to(dtype=return_dtype)
# Diffusers FreeU
original_fft_ifftn = torch.fft.ifftn
@wraps(torch.fft.ifftn)
@@ -131,6 +130,7 @@ def from_numpy(ndarray):
else:
return original_from_numpy(ndarray)
original_as_tensor = torch.as_tensor
@wraps(torch.as_tensor)
def as_tensor(data, dtype=None, device=None):
@@ -156,6 +156,7 @@ def torch_tensor(data, *args, dtype=None, device=None, **kwargs):
dtype = torch.float32
return original_torch_tensor(data, *args, dtype=dtype, device=device, **kwargs)
torch.Tensor.original_Tensor_to = torch.Tensor.to
@wraps(torch.Tensor.to)
def Tensor_to(self, device=None, *args, **kwargs):
@@ -169,6 +170,7 @@ def Tensor_to(self, device=None, *args, **kwargs):
device = torch.float32
return self.original_Tensor_to(device, *args, **kwargs)
original_Tensor_cuda = torch.Tensor.cuda
@wraps(torch.Tensor.cuda)
def Tensor_cuda(self, device=None, *args, **kwargs):
@@ -177,6 +179,7 @@ def Tensor_cuda(self, device=None, *args, **kwargs):
else:
return original_Tensor_cuda(self, device, *args, **kwargs)
original_Tensor_pin_memory = torch.Tensor.pin_memory
@wraps(torch.Tensor.pin_memory)
def Tensor_pin_memory(self, device=None, *args, **kwargs):
@@ -185,6 +188,7 @@ def Tensor_pin_memory(self, device=None, *args, **kwargs):
else:
return original_Tensor_pin_memory(self, device, *args, **kwargs)
original_UntypedStorage_init = torch.UntypedStorage.__init__
@wraps(torch.UntypedStorage.__init__)
def UntypedStorage_init(*args, device=None, **kwargs):
@@ -193,6 +197,7 @@ def UntypedStorage_init(*args, device=None, **kwargs):
else:
return original_UntypedStorage_init(*args, device=device, **kwargs)
if torch_version[0] > 2 or (torch_version[0] == 2 and torch_version[1] >= 4):
original_UntypedStorage_to = torch.UntypedStorage.to
@wraps(torch.UntypedStorage.to)
@@ -210,6 +215,7 @@ if torch_version[0] > 2 or (torch_version[0] == 2 and torch_version[1] >= 4):
else:
return original_UntypedStorage_cuda(self, device=device, non_blocking=non_blocking, **kwargs)
original_torch_empty = torch.empty
@wraps(torch.empty)
def torch_empty(*args, device=None, **kwargs):
@@ -218,6 +224,7 @@ def torch_empty(*args, device=None, **kwargs):
else:
return original_torch_empty(*args, device=device, **kwargs)
original_torch_randn = torch.randn
@wraps(torch.randn)
def torch_randn(*args, device=None, dtype=None, **kwargs):
@@ -228,6 +235,7 @@ def torch_randn(*args, device=None, dtype=None, **kwargs):
else:
return original_torch_randn(*args, device=device, dtype=dtype, **kwargs)
original_torch_ones = torch.ones
@wraps(torch.ones)
def torch_ones(*args, device=None, **kwargs):
@@ -236,6 +244,7 @@ def torch_ones(*args, device=None, **kwargs):
else:
return original_torch_ones(*args, device=device, **kwargs)
original_torch_zeros = torch.zeros
@wraps(torch.zeros)
def torch_zeros(*args, device=None, **kwargs):
@@ -244,6 +253,7 @@ def torch_zeros(*args, device=None, **kwargs):
else:
return original_torch_zeros(*args, device=device, **kwargs)
original_torch_full = torch.full
@wraps(torch.full)
def torch_full(*args, device=None, **kwargs):
@@ -252,6 +262,7 @@ def torch_full(*args, device=None, **kwargs):
else:
return original_torch_full(*args, device=device, **kwargs)
original_torch_linspace = torch.linspace
@wraps(torch.linspace)
def torch_linspace(*args, device=None, **kwargs):
@@ -260,6 +271,7 @@ def torch_linspace(*args, device=None, **kwargs):
else:
return original_torch_linspace(*args, device=device, **kwargs)
original_torch_eye = torch.eye
@wraps(torch.eye)
def torch_eye(*args, device=None, **kwargs):
@@ -268,6 +280,7 @@ def torch_eye(*args, device=None, **kwargs):
else:
return original_torch_eye(*args, device=device, **kwargs)
original_torch_load = torch.load
@wraps(torch.load)
def torch_load(f, map_location=None, *args, **kwargs):
@@ -276,6 +289,7 @@ def torch_load(f, map_location=None, *args, **kwargs):
else:
return original_torch_load(f, *args, map_location=map_location, **kwargs)
@wraps(torch.cuda.synchronize)
def torch_cuda_synchronize(device=None):
if check_cuda(device):
@@ -283,6 +297,7 @@ def torch_cuda_synchronize(device=None):
else:
return torch.xpu.synchronize(device)
@wraps(torch.cuda.device)
def torch_cuda_device(device):
if check_cuda(device):
@@ -290,6 +305,7 @@ def torch_cuda_device(device):
else:
return torch.xpu.device(device)
@wraps(torch.cuda.set_device)
def torch_cuda_set_device(device):
if check_cuda(device):
@@ -297,6 +313,7 @@ def torch_cuda_set_device(device):
else:
torch.xpu.set_device(device)
# torch.Generator has to be a class for isinstance checks
original_torch_Generator = torch.Generator
class torch_Generator(original_torch_Generator):
@@ -335,7 +352,6 @@ def ipex_hijacks():
torch._C.Generator = torch_Generator
torch.backends.cuda.sdp_kernel = return_null_context
torch.nn.DataParallel = DummyDataParallel
torch.UntypedStorage.is_cuda = is_cuda
torch.amp.autocast_mode.autocast.__init__ = autocast_init