From b075d3c8fdf6dece1c66d901a04a651509b2e7fd Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 30 Apr 2023 15:13:56 +0300 Subject: [PATCH] Intel ARC Support --- modules/cmd_args.py | 1 + modules/codeformer_model.py | 6 +- modules/devices.py | 33 ++++++++--- modules/memmon.py | 55 ++++++++++++++----- modules/processing.py | 12 +++- modules/sd_hijack_optimizations.py | 55 ++++++++++++++----- modules/sd_hijack_unet.py | 3 +- modules/sd_models.py | 1 - modules/shared.py | 4 +- .../textual_inversion/textual_inversion.py | 6 +- setup.py | 18 +++++- 11 files changed, 153 insertions(+), 41 deletions(-) diff --git a/modules/cmd_args.py b/modules/cmd_args.py index 9b8438800..fdc9ab240 100644 --- a/modules/cmd_args.py +++ b/modules/cmd_args.py @@ -23,6 +23,7 @@ parser.add_argument("--allow-code", action='store_true', help="Allow custom scri parser.add_argument("--share", action='store_true', help="Enable to make the UI accessible through Gradio site") parser.add_argument("--enable-insecure", action='store_true', help="Enable extensions tab regardless of other options") parser.add_argument("--use-cpu", nargs='+', help="Force use CPU for specified modules", default=[], type=str.lower) +parser.add_argument("--use-ipex", action='store_true', help="Force use Intel OneAPI XPU backend") parser.add_argument("--listen", action='store_true', help="Launch web server using public IP address") parser.add_argument("--port", type=int, help="Launch web server with given server port", default=None) parser.add_argument("--hide-ui-dir-config", action='store_true', help="Hide directory configuration from UI", default=False) diff --git a/modules/codeformer_model.py b/modules/codeformer_model.py index cbe06ec1e..5217f69db 100644 --- a/modules/codeformer_model.py +++ b/modules/codeformer_model.py @@ -103,7 +103,11 @@ def setup_model(dirname): output = self.net(cropped_face_t, w=w if w is not None else shared.opts.code_former_weight, adain=True)[0] restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1)) del output - torch.cuda.empty_cache() + from modules import shared + if shared.cmd_opts.use_ipex: + torch.xpu.empty_cache() + else: + torch.cuda.empty_cache() except Exception as error: print(f'\tFailed inference for CodeFormer: {error}', file=sys.stderr) restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1)) diff --git a/modules/devices.py b/modules/devices.py index e317d91f4..3606597d3 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -22,9 +22,13 @@ def extract_device_id(args, name): def get_cuda_device_string(): from modules import shared - if shared.cmd_opts.device_id is not None: - return f"cuda:{shared.cmd_opts.device_id}" - return "cuda" + if shared.cmd_opts.use_ipex: + return "xpu" + else: + from modules import shared + if shared.cmd_opts.device_id is not None: + return f"cuda:{shared.cmd_opts.device_id}" + return "cuda" def get_dml_device_string(): @@ -35,7 +39,10 @@ def get_dml_device_string(): def get_optimal_device_name(): - if torch.cuda.is_available(): + from modules import shared + if shared.cmd_opts.use_ipex: + return "xpu" + elif torch.cuda.is_available(): return get_cuda_device_string() if has_mps(): return "mps" @@ -61,7 +68,11 @@ def get_device_for(task): def torch_gc(): - if torch.cuda.is_available(): + from modules import shared + if shared.cmd_opts.use_ipex: + with torch.xpu.device("xpu"): + torch.xpu.empty_cache() + elif torch.cuda.is_available(): with torch.cuda.device(get_cuda_device_string()): torch.cuda.empty_cache() torch.cuda.ipc_collect() @@ -137,11 +148,19 @@ def autocast(disable=False): return contextlib.nullcontext() if dtype == torch.float32 or shared.cmd_opts.precision == "Full": return contextlib.nullcontext() - return torch.autocast("cuda") + from modules import shared + if shared.cmd_opts.use_ipex: + return torch.xpu.amp.autocast(enabled=True, dtype=dtype, cache_enabled=False) + else: + 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() + from modules import shared + if shared.cmd_opts.use_ipex: + return torch.autocast("xpu", enabled=False) if torch.is_autocast_enabled() and not disable else contextlib.nullcontext() + else: + return torch.autocast("cuda", enabled=False) if torch.is_autocast_enabled() and not disable else contextlib.nullcontext() class NansException(Exception): diff --git a/modules/memmon.py b/modules/memmon.py index 9b013e6b4..4ceb29a37 100644 --- a/modules/memmon.py +++ b/modules/memmon.py @@ -19,26 +19,44 @@ class MemUsageMonitor(threading.Thread): self.daemon = True self.run_flag = threading.Event() self.data = defaultdict(int) - if not torch.cuda.is_available(): + from modules import shared + if not torch.cuda.is_available() or not shared.cmd_opts.use_ipex: self.disabled = True else: - try: - self.cuda_mem_get_info() - torch.cuda.memory_stats(self.device) - except Exception as e: # AMD or whatever - print(f"Torch exception: {e}") - self.disabled = True + if shared.cmd_opts.use_ipex: + try: + self.cuda_mem_get_info() + torch.cuda.memory_stats("xpu") + except Exception as e: # AMD or whatever + print(f"Torch exception: {e}") + self.disabled = True + + else: + try: + self.cuda_mem_get_info() + torch.cuda.memory_stats(self.device) + except Exception as e: # AMD or whatever + print(f"Torch exception: {e}") + self.disabled = True def cuda_mem_get_info(self): - index = self.device.index if self.device.index is not None else torch.cuda.current_device() - return torch.cuda.mem_get_info(index) + from modules import shared + if shared.cmd_opts.use_ipex: + return torch.xpu.mem_get_info("xpu") + else: + index = self.device.index if self.device.index is not None else torch.cuda.current_device() + return torch.cuda.mem_get_info(index) def run(self): if self.disabled: return while True: self.run_flag.wait() - torch.cuda.reset_peak_memory_stats() + from modules import shared + if shared.cmd_opts.use_ipex: + torch.xpu.reset_peak_memory_stats() + else: + torch.cuda.reset_peak_memory_stats() self.data.clear() if self.opts.memmon_poll_rate <= 0: self.run_flag.clear() @@ -54,12 +72,19 @@ class MemUsageMonitor(threading.Thread): for k, v in self.read().items(): print(k, -(v // -(1024 ** 2))) print(self, 'raw torch memory stats:') - tm = torch.cuda.memory_stats(self.device) + from modules import shared + if shared.cmd_opts.use_ipex: + tm = torch.xpu.memory_stats("xpu") + else: + tm = torch.cuda.memory_stats(self.device) for k, v in tm.items(): if 'bytes' not in k: continue print('\t' if 'peak' in k else '', k, -(v // -(1024 ** 2))) - print(torch.cuda.memory_summary()) + if shared.cmd_opts.use_ipex: + print(torch.xpu.memory_summary()) + else: + print(torch.cuda.memory_summary()) def monitor(self): self.run_flag.set() @@ -70,7 +95,11 @@ class MemUsageMonitor(threading.Thread): self.data["free"] = free self.data["total"] = total - torch_stats = torch.cuda.memory_stats(self.device) + from modules import shared + if shared.cmd_opts.use_ipex: + torch_stats = torch.xpu.memory_stats("xpu") + else: + torch_stats = torch.cuda.memory_stats(self.device) self.data["active"] = torch_stats["active.all.current"] self.data["active_peak"] = torch_stats["active_bytes.all.peak"] self.data["reserved"] = torch_stats["reserved_bytes.all.current"] diff --git a/modules/processing.py b/modules/processing.py index e793f12a3..293a8d606 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -55,7 +55,17 @@ def memory_stats(): except Exception as e: mem.update({ 'ram': e }) try: - if torch.cuda.is_available(): + from modules import shared + if shared.cmd_opts.use_ipex: + s = torch.xpu.mem_get_info() + gpu = { 'used': gb(s[1] - s[0]), 'total': gb(s[1]) } + s = dict(torch.xpu.memory_stats("xpu")) + mem.update({ + 'gpu': gpu, + 'retries': s['num_alloc_retries'], + 'oom': s['num_ooms'] + }) + elif torch.cuda.is_available(): s = torch.cuda.mem_get_info() gpu = { 'used': gb(s[1] - s[0]), 'total': gb(s[1]) } s = dict(torch.cuda.memory_stats(shared.device)) diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py index 12ee9f956..5168b4b7a 100644 --- a/modules/sd_hijack_optimizations.py +++ b/modules/sd_hijack_optimizations.py @@ -22,7 +22,15 @@ if shared.opts.cross_attention_optimization == "xFormers": def get_available_vram(): - if shared.device.type == 'cuda': + if shared.cmd_opts.use_ipex: + stats = torch.xpu.memory_stats("xpu") + mem_active = stats['active_bytes.all.current'] + mem_reserved = stats['reserved_bytes.all.current'] + mem_free_xpu, _ = torch.xpu.mem_get_info("xpu") + mem_free_torch = mem_reserved - mem_active + mem_free_total = mem_free_xpu + mem_free_torch + return mem_free_total + elif shared.device.type == 'cuda': stats = torch.cuda.memory_stats(shared.device) mem_active = stats['active_bytes.all.current'] mem_reserved = stats['reserved_bytes.all.current'] @@ -189,14 +197,24 @@ def einsum_op_tensor_mem(q, k, v, max_tensor_mb): return einsum_op_slice_1(q, k, v, max(q.shape[1] // div, 1)) def einsum_op_cuda(q, k, v): - stats = torch.cuda.memory_stats(q.device) - mem_active = stats['active_bytes.all.current'] - mem_reserved = stats['reserved_bytes.all.current'] - mem_free_cuda, _ = torch.cuda.mem_get_info(q.device) - mem_free_torch = mem_reserved - mem_active - mem_free_total = mem_free_cuda + mem_free_torch - # Divide factor of safety as there's copying and fragmentation - return einsum_op_tensor_mem(q, k, v, mem_free_total / 3.3 / (1 << 20)) + if shared.cmd_opts.use_ipex: + stats = torch.xpu.memory_stats("xpu") + mem_active = stats['active_bytes.all.current'] + mem_reserved = stats['reserved_bytes.all.current'] + mem_free_xpu, _ = torch.xpu.mem_get_info("xpu") + mem_free_torch = mem_reserved - mem_active + mem_free_total = mem_free_xpu + mem_free_torch + # Divide factor of safety as there's copying and fragmentation + return einsum_op_tensor_mem(q, k, v, mem_free_total / 3.3 / (1 << 20)) + else: + stats = torch.cuda.memory_stats(q.device) + mem_active = stats['active_bytes.all.current'] + mem_reserved = stats['reserved_bytes.all.current'] + mem_free_cuda, _ = torch.cuda.mem_get_info(q.device) + mem_free_torch = mem_reserved - mem_active + mem_free_total = mem_free_cuda + mem_free_torch + # Divide factor of safety as there's copying and fragmentation + return einsum_op_tensor_mem(q, k, v, mem_free_total / 3.3 / (1 << 20)) def einsum_op_dml(q, k, v): mem_total, mem_active = torch.dml.memory_stats(q.device) @@ -204,6 +222,9 @@ def einsum_op_dml(q, k, v): return einsum_op_tensor_mem(q, k, v, (mem_reserved - mem_active) if mem_reserved > mem_active else 1) def einsum_op(q, k, v): + if shared.cmd_opts.use_ipex: + return einsum_op_cuda(q, k, v) + if q.device.type == 'cuda': return einsum_op_cuda(q, k, v) @@ -397,8 +418,12 @@ def scaled_dot_product_attention_forward(self, x, context=None, mask=None): return hidden_states def scaled_dot_product_no_mem_attention_forward(self, x, context=None, mask=None): - with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False): - return scaled_dot_product_attention_forward(self, x, context, mask) + if shared.cmd_opts.use_ipex: + with torch.backends.xpu.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False): + return scaled_dot_product_attention_forward(self, x, context, mask) + else: + with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False): + return scaled_dot_product_attention_forward(self, x, context, mask) def cross_attention_attnblock_forward(self, x): h_ = x @@ -502,8 +527,12 @@ def sdp_attnblock_forward(self, x): return x + out def sdp_no_mem_attnblock_forward(self, x): - with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False): - return sdp_attnblock_forward(self, x) + if shared.cmd_opts.use_ipex: + with torch.backends.xpu.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False): + return sdp_attnblock_forward(self, x) + else: + with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False): + return sdp_attnblock_forward(self, x) def sub_quad_attnblock_forward(self, x): h_ = x diff --git a/modules/sd_hijack_unet.py b/modules/sd_hijack_unet.py index 158582632..7ff553ae3 100644 --- a/modules/sd_hijack_unet.py +++ b/modules/sd_hijack_unet.py @@ -3,6 +3,7 @@ from packaging import version from modules import devices from modules.sd_hijack_utils import CondFunc +from modules import shared class TorchHijackForUnet: @@ -67,7 +68,7 @@ def hijack_ddpm_edit(): unet_needs_upcast = lambda *args, **kwargs: devices.unet_needs_upcast CondFunc('ldm.models.diffusion.ddpm.LatentDiffusion.apply_model', apply_model, unet_needs_upcast) CondFunc('ldm.modules.diffusionmodules.openaimodel.timestep_embedding', lambda orig_func, timesteps, *args, **kwargs: orig_func(timesteps, *args, **kwargs).to(torch.float32 if timesteps.dtype == torch.int64 else devices.dtype_unet), unet_needs_upcast) -if version.parse(torch.__version__) <= version.parse("1.13.2") or torch.cuda.is_available(): +if version.parse(torch.__version__) <= version.parse("1.13.2") or torch.cuda.is_available() or shared.cmd_opts.use_ipex: CondFunc('ldm.modules.diffusionmodules.util.GroupNorm32.forward', lambda orig_func, self, *args, **kwargs: orig_func(self.float(), *args, **kwargs), unet_needs_upcast) CondFunc('ldm.modules.attention.GEGLU.forward', lambda orig_func, self, x: orig_func(self.float(), x.float()).to(devices.dtype_unet), unet_needs_upcast) CondFunc('open_clip.transformer.ResidualAttentionBlock.__init__', lambda orig_func, *args, **kwargs: kwargs.update({'act_layer': GELUHijack}) and False or orig_func(*args, **kwargs), lambda _, *args, **kwargs: kwargs.get('act_layer') is None or kwargs['act_layer'] == torch.nn.GELU) diff --git a/modules/sd_models.py b/modules/sd_models.py index a2ef7a012..49de37097 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -533,7 +533,6 @@ def unload_model_weights(sd_model=None, _info=None): sd_model = None gc.collect() devices.torch_gc() - torch.cuda.empty_cache() print(f"Unloaded weights {timer.summary()}") return sd_model diff --git a/modules/shared.py b/modules/shared.py index f3ebbce55..aa2ede29e 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -238,7 +238,7 @@ options_templates.update(options_section(('sd', "Stable Diffusion"), { "comma_padding_backtrack": OptionInfo(20, "Increase coherency by padding from the last comma within n tokens when using more than 75 tokens", gr.Slider, {"minimum": 0, "maximum": 74, "step": 1 }), "CLIP_stop_at_last_layers": OptionInfo(1, "Clip skip", gr.Slider, {"minimum": 1, "maximum": 12, "step": 1, "visible": False}), "upcast_attn": OptionInfo(False, "Upcast cross attention layer to float32"), - "cross_attention_optimization": OptionInfo("Sub-quadratic" if is_device_dml else "Scaled-Dot-Product", "Cross-attention optimization method", gr.Radio, lambda: {"choices": shared_items.list_crossattention() }), + "cross_attention_optimization": OptionInfo("Sub-quadratic" if is_device_dml else "Split attention" if cmd_opts.use_ipex else "Scaled-Dot-Product", "Cross-attention optimization method", gr.Radio, lambda: {"choices": shared_items.list_crossattention() }), "cross_attention_options": OptionInfo([], "Cross-attention advanced options", gr.CheckboxGroup, lambda: {"choices": ['xFormers enable flash Attention', 'SDP disable memory attention']}), "sub_quad_q_chunk_size": OptionInfo(512, "Sub-quadratic cross-attention query chunk size for the layer optimization to use", gr.Slider, {"minimum": 16, "maximum": 8192, "step": 8}), "sub_quad_kv_chunk_size": OptionInfo(512, "Sub-quadratic cross-attentionkv chunk size for the sub-quadratic cross-attention layer optimization to use", gr.Slider, {"minimum": 0, "maximum": 8192, "step": 8}), @@ -318,7 +318,7 @@ options_templates.update(options_section(('cuda', "CUDA Settings"), { "cuda_dtype": OptionInfo("FP32" if sys.platform == "darwin" else "FP16", "Device precision type", gr.Radio, lambda: {"choices": ["FP32", "FP16", "BF16"]}), "no_half": OptionInfo(True if is_device_dml else False, "Use full precision for model (--no-half)", None, None, lambda: print("Warning: Most of DirectML devices do not fully support half mode. Recommend to use full precision to model.") if is_device_dml else None), "no_half_vae": OptionInfo(True if is_device_dml else False, "Use full precision for VAE (--no-half-vae)"), - "upcast_sampling": OptionInfo(True if sys.platform == "darwin" else False, "Enable upcast sampling. Usually produces similar results to --no-half with better performance while using less memory"), + "upcast_sampling": OptionInfo(True if sys.platform == "darwin" or cmd_opts.use_ipex else False, "Enable upcast sampling. Usually produces similar results to --no-half with better performance while using less memory"), "disable_nan_check": OptionInfo(True, "Do not check if produced images/latent spaces have NaN values"), "rollback_vae": OptionInfo(False, "Attempt to roll back VAE when produced NaN values, requires NaN check (experimental)"), "opt_channelslast": OptionInfo(False, "Use channels last as torch memory format "), diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py index 2180d7f32..36a1e1e17 100644 --- a/modules/textual_inversion/textual_inversion.py +++ b/modules/textual_inversion/textual_inversion.py @@ -434,7 +434,11 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st else: print("No saved optimizer exists in checkpoint") - scaler = torch.cuda.amp.GradScaler() + from modules import shared + if shared.cmd_opts.use_ipex: + scaler = torch.xpu.amp.GradScaler() + else: + scaler = torch.cuda.amp.GradScaler() batch_size = ds.batch_size gradient_step = ds.gradient_step diff --git a/setup.py b/setup.py index cfaa1aaa0..59c200161 100644 --- a/setup.py +++ b/setup.py @@ -56,6 +56,7 @@ def setup_logging(clean=False): # check if package is installed def installed(package, friendly: str = None): import pkg_resources + from modules import shared ok = True try: if friendly: @@ -76,6 +77,8 @@ def installed(package, friendly: str = None): ok = ok and spec is not None if ok: version = pkg_resources.get_distribution(p[0]).version + if shared.cmd_opts.use_ipex and p[0] == "pytorch_lightning": + p[1] = "1.8.6" log.debug(f"Package version found: {p[0]} {version}") if len(p) > 1: ok = ok and version == p[1] @@ -91,6 +94,9 @@ def installed(package, friendly: str = None): # install package using pip if not already installed def install(package, friendly: str = None, ignore: bool = False): + from modules import shared + if shared.cmd_opts.use_ipex and package == "pytorch_lightning==1.9.4": + package = "pytorch_lightning==1.8.6" def pip(arg: str): arg = arg.replace('>=', '==') log.info(f'Installing package: {arg.replace("install", "").replace("--upgrade", "").replace("--no-deps", "").replace(" ", " ").strip()}') @@ -188,6 +194,7 @@ def check_python(): # check torch version def check_torch(): + from modules import shared if shutil.which('nvidia-smi') is not None or os.path.exists(os.path.join(os.environ.get('SystemRoot') or r'C:\Windows', 'System32', 'nvidia-smi.exe')): log.info('nVidia toolkit detected') torch_command = os.environ.get('TORCH_COMMAND', 'torch torchaudio torchvision --index-url https://download.pytorch.org/whl/cu118') @@ -197,6 +204,11 @@ def check_torch(): os.environ.setdefault('HSA_OVERRIDE_GFX_VERSION', '10.3.0') torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm5.4.2') xformers_package = os.environ.get('XFORMERS_PACKAGE', 'none') + elif shutil.which('sycl-ls') is not None or os.path.exists('/opt/intel/oneapi'): + shared.cmd_opts.use_ipex = True + log.info('Intel toolkit detected') + torch_command = os.environ.get('TORCH_COMMAND', 'torch==1.13.0a0+git6c9b55e torchvision==0.14.1a0 intel_extension_for_pytorch==1.13.120+xpu --index-url https://developer.intel.com/ipex-whl-stable-xpu') + xformers_package = os.environ.get('XFORMERS_PACKAGE', 'none') else: machine = platform.machine() if 'arm' not in machine and 'aarch' not in machine and not args.nodirectml: # torch-directml is available on AMD64 @@ -212,7 +224,11 @@ def check_torch(): try: import torch log.info(f'Torch {torch.__version__}') - if torch.cuda.is_available(): + if shared.cmd_opts.use_ipex: + import intel_extension_for_pytorch as ipex + log.info(f'Torch backend: Intel OneAPI {torch.__version__}') + log.info(f'Torch detected GPU: {torch.xpu.get_device_name("xpu")} VRAM {round(torch.xpu.get_device_properties("xpu").total_memory / 1024 / 1024)}') + elif torch.cuda.is_available(): if torch.version.cuda: log.info(f'Torch backend: nVidia CUDA {torch.version.cuda} cuDNN {torch.backends.cudnn.version() if torch.backends.cudnn.is_available() else "N/A"}') elif torch.version.hip: