diff --git a/extensions-builtin/sd-webui-controlnet b/extensions-builtin/sd-webui-controlnet index d2da774a4..af4720780 160000 --- a/extensions-builtin/sd-webui-controlnet +++ b/extensions-builtin/sd-webui-controlnet @@ -1 +1 @@ -Subproject commit d2da774a40ff9c3770e21f71fb516403022fc3f6 +Subproject commit af4720780f10d912789cbd6db1fbc6d2f0afc533 diff --git a/modules/api/api.py b/modules/api/api.py index 0717edfaf..fdd26f868 100644 --- a/modules/api/api.py +++ b/modules/api/api.py @@ -573,7 +573,24 @@ class Api: ram = { 'error': f'{err}' } try: import torch - if torch.cuda.is_available(): + if shared.cmd_opts.use_ipex(): + import intel_extension_for_pytorch as ipex + system = { 'free': (torch.xpu.get_device_properties("xpu").total_memory - torch.xpu.memory_allocated()), 'used': torch.xpu.memory_allocated(), 'total': torch.xpu.get_device_properties("xpu").total_memory } + s = dict(torch.xpu.memory_stats("xpu")) + allocated = { 'current': s['allocated_bytes.all.current'], 'peak': s['allocated_bytes.all.peak'] } + reserved = { 'current': s['reserved_bytes.all.current'], 'peak': s['reserved_bytes.all.peak'] } + active = { 'current': s['active_bytes.all.current'], 'peak': s['active_bytes.all.peak'] } + inactive = { 'current': s['inactive_split_bytes.all.current'], 'peak': s['inactive_split_bytes.all.peak'] } + warnings = { 'retries': s['num_alloc_retries'], 'oom': s['num_ooms'] } + cuda = { + 'system': system, + 'active': active, + 'allocated': allocated, + 'reserved': reserved, + 'inactive': inactive, + 'events': warnings, + } + elif torch.cuda.is_available(): s = torch.cuda.mem_get_info() system = { 'free': s[0], 'used': s[1] - s[0], 'total': s[1] } s = dict(torch.cuda.memory_stats(shared.device)) diff --git a/modules/cmd_args.py b/modules/cmd_args.py index c45e6c203..0d1da94ef 100644 --- a/modules/cmd_args.py +++ b/modules/cmd_args.py @@ -27,6 +27,7 @@ group.add_argument("--allow-code", action='store_true', help="Allow custom scrip group.add_argument("--share", action='store_true', help="Enable UI accessible through Gradio site, default: %(default)s") group.add_argument("--insecure", action='store_true', help="Enable extensions tab regardless of other options, default: %(default)s") group.add_argument("--use-cpu", nargs='+', default=[], type=str.lower, help="Force use CPU for specified modules, default: %(default)s") +group.add_argument("--use-ipex", action='store_true', help="Force use Intel OneAPI XPU backend, default: %(default)s", default=False) group.add_argument("--listen", action='store_true', help="Launch web server using public IP address, default: %(default)s") group.add_argument("--port", type=int, default=7860, help="Launch web server with given server port, default: %(default)s") group.add_argument("--freeze", action='store_true', help="Disable editing settings", default=False) diff --git a/modules/codeformer/codeformer_arch.py b/modules/codeformer/codeformer_arch.py index 11dcc3ee7..6d7b926fe 100644 --- a/modules/codeformer/codeformer_arch.py +++ b/modules/codeformer/codeformer_arch.py @@ -3,6 +3,10 @@ import math import numpy as np import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass from torch import nn, Tensor import torch.nn.functional as F from typing import Optional, List diff --git a/modules/codeformer/vqgan_arch.py b/modules/codeformer/vqgan_arch.py index e72936838..e66bb2a72 100644 --- a/modules/codeformer/vqgan_arch.py +++ b/modules/codeformer/vqgan_arch.py @@ -7,6 +7,10 @@ https://github.com/samb-t/unleashing-transformers/blob/master/models/vqgan.py ''' import numpy as np import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass import torch.nn as nn import torch.nn.functional as F import copy diff --git a/modules/codeformer_model.py b/modules/codeformer_model.py index cbe06ec1e..9d75e823d 100644 --- a/modules/codeformer_model.py +++ b/modules/codeformer_model.py @@ -3,6 +3,10 @@ import sys import cv2 import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass import modules.face_restoration from modules import shared, devices, modelloader, errors @@ -103,7 +107,10 @@ 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() + if 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/deepbooru.py b/modules/deepbooru.py index 1c4554a20..50e400fd8 100644 --- a/modules/deepbooru.py +++ b/modules/deepbooru.py @@ -2,6 +2,10 @@ import os import re import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass import numpy as np from modules import modelloader, paths, deepbooru_model, devices, images, shared diff --git a/modules/deepbooru_model.py b/modules/deepbooru_model.py index c2c77cd25..ef53494a2 100644 --- a/modules/deepbooru_model.py +++ b/modules/deepbooru_model.py @@ -1,4 +1,8 @@ import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass import torch.nn as nn import torch.nn.functional as F diff --git a/modules/devices.py b/modules/devices.py index e317d91f4..8be1e3866 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -1,6 +1,11 @@ import sys import contextlib import torch +from modules import shared +try: + import intel_extension_for_pytorch as ipex +except: + pass if sys.platform == "darwin": from modules import mac_specific @@ -21,21 +26,24 @@ 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: + if shared.cmd_opts.device_id is not None: + return f"cuda:{shared.cmd_opts.device_id}" + return "cuda" def get_dml_device_string(): - from modules import shared if shared.cmd_opts.device_id is not None: return f"privateuseone:{shared.cmd_opts.device_id}" return "privateuseone:0" def get_optimal_device_name(): - if torch.cuda.is_available(): + if shared.cmd_opts.use_ipex: + return "xpu" + elif torch.cuda.is_available(): return get_cuda_device_string() if has_mps(): return "mps" @@ -54,21 +62,22 @@ def get_optimal_device(): def get_device_for(task): - from modules import shared if task in shared.cmd_opts.use_cpu: return cpu return get_optimal_device() def torch_gc(): - if torch.cuda.is_available(): + 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() def set_cuda_params(): - from modules import shared if torch.cuda.is_available(): try: torch.backends.cuda.matmul.allow_tf32 = shared.opts.cuda_allow_tf32 @@ -132,16 +141,21 @@ def randn_without_seed(shape): def autocast(disable=False): - from modules import shared if disable: return contextlib.nullcontext() if dtype == torch.float32 or shared.cmd_opts.precision == "Full": return contextlib.nullcontext() - return torch.autocast("cuda") + 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() + 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): @@ -149,7 +163,6 @@ class NansException(Exception): def test_for_nans(x, where): - from modules import shared if shared.opts.disable_nan_check: return if not torch.all(torch.isnan(x)).item(): diff --git a/modules/esrgan_model.py b/modules/esrgan_model.py index bb4c6619b..769d66f01 100644 --- a/modules/esrgan_model.py +++ b/modules/esrgan_model.py @@ -2,6 +2,10 @@ import os import numpy as np import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass from PIL import Image from basicsr.utils.download_util import load_file_from_url diff --git a/modules/esrgan_model_arch.py b/modules/esrgan_model_arch.py index 411d98d38..fc352d0ba 100644 --- a/modules/esrgan_model_arch.py +++ b/modules/esrgan_model_arch.py @@ -2,6 +2,10 @@ import math import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass import torch.nn as nn import torch.nn.functional as F diff --git a/modules/extras.py b/modules/extras.py index c0ae9477f..4513f2491 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -4,6 +4,10 @@ import html import shutil import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass import tqdm import gradio as gr import safetensors.torch diff --git a/modules/hypernetworks/hypernetwork.py b/modules/hypernetworks/hypernetwork.py index 4aa5ffcdc..a1caecbe4 100644 --- a/modules/hypernetworks/hypernetwork.py +++ b/modules/hypernetworks/hypernetwork.py @@ -8,6 +8,10 @@ import inspect import modules.textual_inversion.dataset import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass import tqdm from einops import rearrange, repeat from ldm.util import default @@ -591,7 +595,10 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi print("Cannot resume from saved optimizer!") print(e) - scaler = torch.cuda.amp.GradScaler() + 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 @@ -708,7 +715,9 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi hypernetwork.eval() rng_state = torch.get_rng_state() cuda_rng_state = None - if torch.cuda.is_available(): + if shared.cmd_opts.use_ipex: + cuda_rng_state = torch.xpu.get_rng_state_all() + elif torch.cuda.is_available(): cuda_rng_state = torch.cuda.get_rng_state_all() shared.sd_model.cond_stage_model.to(devices.device) shared.sd_model.first_stage_model.to(devices.device) @@ -745,7 +754,9 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi shared.sd_model.cond_stage_model.to(devices.cpu) shared.sd_model.first_stage_model.to(devices.cpu) torch.set_rng_state(rng_state) - if torch.cuda.is_available(): + if shared.cmd_opts.use_ipex: + torch.xpu.set_rng_state_all(cuda_rng_state) + elif torch.cuda.is_available(): torch.cuda.set_rng_state_all(cuda_rng_state) hypernetwork.train() if image is not None: diff --git a/modules/interrogate.py b/modules/interrogate.py index 6afbde570..93bb08f20 100644 --- a/modules/interrogate.py +++ b/modules/interrogate.py @@ -5,6 +5,10 @@ from pathlib import Path import re import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass import torch.hub from torchvision import transforms diff --git a/modules/lowvram.py b/modules/lowvram.py index e254cc131..7dba01593 100644 --- a/modules/lowvram.py +++ b/modules/lowvram.py @@ -1,4 +1,8 @@ import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass from modules import devices module_in_gpu = None diff --git a/modules/mac_specific.py b/modules/mac_specific.py index c8a534d0e..2455800d5 100644 --- a/modules/mac_specific.py +++ b/modules/mac_specific.py @@ -1,4 +1,8 @@ import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass import platform from modules.sd_hijack_utils import CondFunc from packaging import version diff --git a/modules/memmon.py b/modules/memmon.py index 9b013e6b4..3abc70ac3 100644 --- a/modules/memmon.py +++ b/modules/memmon.py @@ -2,6 +2,12 @@ import threading import time from collections import defaultdict import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass + +from modules import shared class MemUsageMonitor(threading.Thread): @@ -19,26 +25,41 @@ class MemUsageMonitor(threading.Thread): self.daemon = True self.run_flag = threading.Event() self.data = defaultdict(int) - if not torch.cuda.is_available(): + 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) + if shared.cmd_opts.use_ipex: + return [(torch.xpu.get_device_properties("xpu").total_memory - torch.xpu.memory_allocated()), torch.xpu.get_device_properties("xpu").total_memory] + 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() + 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 +75,18 @@ 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) + 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 +97,10 @@ class MemUsageMonitor(threading.Thread): self.data["free"] = free self.data["total"] = total - torch_stats = torch.cuda.memory_stats(self.device) + 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/models/diffusion/ddpm_edit.py b/modules/models/diffusion/ddpm_edit.py index f3d49c44c..846a74fc4 100644 --- a/modules/models/diffusion/ddpm_edit.py +++ b/modules/models/diffusion/ddpm_edit.py @@ -10,6 +10,10 @@ https://github.com/CompVis/taming-transformers # See more details in LICENSE. import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass import torch.nn as nn import numpy as np import pytorch_lightning as pl diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py index 3100522ab..6dd7c7fd8 100644 --- a/modules/models/diffusion/uni_pc/sampler.py +++ b/modules/models/diffusion/uni_pc/sampler.py @@ -2,6 +2,10 @@ import numpy as np import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass from .uni_pc import NoiseScheduleVP, model_wrapper, UniPC from modules import shared, devices diff --git a/modules/models/diffusion/uni_pc/uni_pc.py b/modules/models/diffusion/uni_pc/uni_pc.py index 61ee39522..895fc58c3 100644 --- a/modules/models/diffusion/uni_pc/uni_pc.py +++ b/modules/models/diffusion/uni_pc/uni_pc.py @@ -1,4 +1,8 @@ import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass import torch.nn.functional as F import math import time diff --git a/modules/processing.py b/modules/processing.py index e793f12a3..36737fdbe 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -8,6 +8,10 @@ from typing import Any, Dict, List import psutil import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass import numpy as np from PIL import Image, ImageFilter, ImageOps import cv2 @@ -55,7 +59,15 @@ def memory_stats(): except Exception as e: mem.update({ 'ram': e }) try: - if torch.cuda.is_available(): + if cmd_opts.use_ipex: + gpu = { 'used': gb(torch.xpu.memory_allocated()), 'total': gb(torch.xpu.get_device_properties("xpu").total_memory) } + 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/prompt_parser.py b/modules/prompt_parser.py index 7006f2822..6722d9f80 100644 --- a/modules/prompt_parser.py +++ b/modules/prompt_parser.py @@ -368,3 +368,7 @@ if __name__ == "__main__": doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE) else: import torch # doctest faster + try: + import intel_extension_for_pytorch as ipex + except: + pass diff --git a/modules/safe.py b/modules/safe.py index 9a1133ddc..dd463ccdd 100644 --- a/modules/safe.py +++ b/modules/safe.py @@ -6,6 +6,10 @@ import zipfile import re import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass import numpy import _codecs diff --git a/modules/sd_disable_initialization.py b/modules/sd_disable_initialization.py index c4a09d15d..5cc5e4e7a 100644 --- a/modules/sd_disable_initialization.py +++ b/modules/sd_disable_initialization.py @@ -1,6 +1,10 @@ import ldm.modules.encoders.modules import open_clip import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass import transformers.utils.hub diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py index 459dfd091..49acdf3dd 100644 --- a/modules/sd_hijack.py +++ b/modules/sd_hijack.py @@ -1,6 +1,10 @@ from types import MethodType from rich import print # pylint: disable=redefined-builtin import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass from torch.nn.functional import silu import ldm.modules.attention import ldm.modules.diffusionmodules.model diff --git a/modules/sd_hijack_clip.py b/modules/sd_hijack_clip.py index 945f7732d..cf4abf84f 100644 --- a/modules/sd_hijack_clip.py +++ b/modules/sd_hijack_clip.py @@ -2,6 +2,10 @@ import math from collections import namedtuple import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass from modules import prompt_parser, devices, sd_hijack from modules.shared import opts diff --git a/modules/sd_hijack_inpainting.py b/modules/sd_hijack_inpainting.py index 4b23c132d..1a9ea9b4c 100644 --- a/modules/sd_hijack_inpainting.py +++ b/modules/sd_hijack_inpainting.py @@ -1,4 +1,8 @@ import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass import ldm.models.diffusion.ddpm import ldm.models.diffusion.ddim diff --git a/modules/sd_hijack_open_clip.py b/modules/sd_hijack_open_clip.py index f76fc1f3b..c0c204a82 100644 --- a/modules/sd_hijack_open_clip.py +++ b/modules/sd_hijack_open_clip.py @@ -1,5 +1,9 @@ import open_clip.tokenizer import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass from modules import sd_hijack_clip, devices diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py index 12ee9f956..3887e238d 100644 --- a/modules/sd_hijack_optimizations.py +++ b/modules/sd_hijack_optimizations.py @@ -2,6 +2,10 @@ import math import psutil import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass from torch import einsum from ldm.util import default @@ -22,7 +26,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.get_device_properties("xpu").total_memory - torch.xpu.memory_allocated()), torch.xpu.get_device_properties("xpu").total_memory] + 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 +201,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.get_device_properties("xpu").total_memory - torch.xpu.memory_allocated()), torch.xpu.get_device_properties("xpu").total_memory] + 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 +226,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 +422,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 +531,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..ce6ac1306 100644 --- a/modules/sd_hijack_unet.py +++ b/modules/sd_hijack_unet.py @@ -1,8 +1,13 @@ import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass from packaging import version from modules import devices from modules.sd_hijack_utils import CondFunc +from modules import shared class TorchHijackForUnet: @@ -67,7 +72,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_hijack_xlmr.py b/modules/sd_hijack_xlmr.py index 28528329b..a9cb9454c 100644 --- a/modules/sd_hijack_xlmr.py +++ b/modules/sd_hijack_xlmr.py @@ -1,4 +1,8 @@ import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass from modules import sd_hijack_clip, devices diff --git a/modules/sd_models.py b/modules/sd_models.py index 4f5b18891..a7a9adeb9 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -8,6 +8,10 @@ from os import mkdir from urllib import request from rich import print, progress # pylint: disable=redefined-builtin import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass import safetensors.torch from omegaconf import OmegaConf import tomesd @@ -533,7 +537,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/sd_models_config.py b/modules/sd_models_config.py index a9c515b14..5bc3799a0 100644 --- a/modules/sd_models_config.py +++ b/modules/sd_models_config.py @@ -1,6 +1,10 @@ import os import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass from modules import paths, sd_disable_initialization diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py index 888f9a30e..dfb478251 100644 --- a/modules/sd_samplers_common.py +++ b/modules/sd_samplers_common.py @@ -1,6 +1,10 @@ from collections import namedtuple import numpy as np import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass from PIL import Image from modules import devices, processing, images, sd_vae_approx diff --git a/modules/sd_samplers_compvis.py b/modules/sd_samplers_compvis.py index 8de719323..6f08a9022 100644 --- a/modules/sd_samplers_compvis.py +++ b/modules/sd_samplers_compvis.py @@ -4,6 +4,10 @@ import ldm.models.diffusion.plms import numpy as np import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass from modules.shared import state from modules import sd_samplers_common, prompt_parser, shared diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py index a30d351fc..5ba34cc33 100644 --- a/modules/sd_samplers_kdiffusion.py +++ b/modules/sd_samplers_kdiffusion.py @@ -1,6 +1,10 @@ from collections import deque import inspect import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass import k_diffusion.sampling from modules import prompt_parser, devices, sd_samplers_common diff --git a/modules/sd_vae.py b/modules/sd_vae.py index e5c544487..a13d73be7 100644 --- a/modules/sd_vae.py +++ b/modules/sd_vae.py @@ -3,8 +3,14 @@ import collections import glob from copy import deepcopy from rich import print # pylint: disable=redefined-builtin +from modules import shared import torch -from modules import paths, shared, devices, script_callbacks, sd_models +try: + import intel_extension_for_pytorch as ipex +except: + if shared.cmd_opts.use_ipex: + print("Failed to import IPEX") +from modules import paths, devices, script_callbacks, sd_models vae_ignore_keys = {"model_ema.decay", "model_ema.num_updates"} diff --git a/modules/sd_vae_approx.py b/modules/sd_vae_approx.py index e2f004683..56c3fb15f 100644 --- a/modules/sd_vae_approx.py +++ b/modules/sd_vae_approx.py @@ -1,6 +1,10 @@ import os import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass from torch import nn from modules import devices, paths diff --git a/modules/shared.py b/modules/shared.py index aabed166c..11177d0ac 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -316,7 +316,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/sub_quadratic_attention.py b/modules/sub_quadratic_attention.py index 87c18a38d..0af680de2 100644 --- a/modules/sub_quadratic_attention.py +++ b/modules/sub_quadratic_attention.py @@ -14,6 +14,10 @@ from functools import partial import math from typing import Optional, NamedTuple, List import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass from torch import Tensor from torch.utils.checkpoint import checkpoint diff --git a/modules/textual_inversion/dataset.py b/modules/textual_inversion/dataset.py index af9fbcf28..272ae76ea 100644 --- a/modules/textual_inversion/dataset.py +++ b/modules/textual_inversion/dataset.py @@ -2,6 +2,10 @@ import os import numpy as np import PIL import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass from PIL import Image from torch.utils.data import Dataset, DataLoader, Sampler from torchvision import transforms diff --git a/modules/textual_inversion/image_embedding.py b/modules/textual_inversion/image_embedding.py index 0ba5db8a4..a2c518af3 100644 --- a/modules/textual_inversion/image_embedding.py +++ b/modules/textual_inversion/image_embedding.py @@ -4,6 +4,10 @@ import numpy as np import zlib from PIL import Image, PngImagePlugin, ImageDraw, ImageFont import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass from modules.shared import opts diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py index cbacc2ce2..ab412f458 100644 --- a/modules/textual_inversion/textual_inversion.py +++ b/modules/textual_inversion/textual_inversion.py @@ -3,6 +3,10 @@ import html import csv from collections import namedtuple import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass import tqdm import safetensors.torch from rich import print # pylint: disable=redefined-builtin @@ -435,7 +439,10 @@ 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() + 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/modules/xlmr.py b/modules/xlmr.py index 9da3161cc..a891beb6d 100644 --- a/modules/xlmr.py +++ b/modules/xlmr.py @@ -1,5 +1,9 @@ from typing import Optional import torch +try: + import intel_extension_for_pytorch as ipex +except: + pass import torch.nn as nn from transformers import XLMRobertaModel,XLMRobertaTokenizer, BertPreTrainedModel, BertModel, BertConfig # pylint: disable=unused-import from transformers.models.xlm_roberta.configuration_xlm_roberta import XLMRobertaConfig diff --git a/setup.py b/setup.py index 3d680f397..393663088 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()}') @@ -187,6 +193,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') @@ -196,6 +203,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.no_directml: # torch-directml is available on AMD64 @@ -211,7 +223,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: diff --git a/webui.py b/webui.py index 6d78be03c..707585ca7 100644 --- a/webui.py +++ b/webui.py @@ -12,6 +12,10 @@ from modules import timer, errors startup_timer = timer.Timer() import torch # pylint: disable=C0411 +try: + import intel_extension_for_pytorch as ipex +except: + pass import torchvision # pylint: disable=W0611,C0411 import pytorch_lightning # pytorch_lightning should be imported after torch, but it re-enables warnings on import so import once to disable them # pylint: disable=W0611,C0411 logging.getLogger("xformers").addFilter(lambda record: 'A matching Triton is not available' not in record.getMessage())