diff --git a/modules/intel/ipex/hijacks.py b/modules/intel/ipex/hijacks.py index e52327971..da24dcbcb 100644 --- a/modules/intel/ipex/hijacks.py +++ b/modules/intel/ipex/hijacks.py @@ -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