Initial IPEX support

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

    if has_mps():

        return "mps"

    if xpu_specific.has_ipex:

        return xpu_specific.get_xpu_device_string()

    return "cpu"

    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):

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)