IPEX rework

installer.py

@@ -315,7 +315,7 @@ def check_torch():
             log.error('Intel OneAPI Toolkit is not activated! Start the WebUI with --use-ipex or activate OneAPI manually')
         os.environ.setdefault('NEOReadDebugKeys', '1')
         os.environ.setdefault('ClDeviceGlobalMemSizeAvailablePercent', '100')
-        torch_command = os.environ.get('TORCH_COMMAND', 'torch==1.13.0a0 torchvision==0.14.1a0 intel_extension_for_pytorch==1.13.120+xpu -f https://developer.intel.com/ipex-whl-stable-xpu')
+        torch_command = os.environ.get('TORCH_COMMAND', 'torch==1.13.0a0+git6c9b55e torchvision==0.14.1a0 intel_extension_for_pytorch==1.13.120+xpu -f https://developer.intel.com/ipex-whl-stable-xpu')
     else:
         machine = platform.machine()
         if sys.platform == 'darwin':

modules/api/api.py

@@ -603,23 +603,7 @@ class Api:
             ram = { 'error': f'{err}' }
         try:
             import torch
-            if devices.backend == 'ipex':
-                system = { 'free': (torch.xpu.get_device_properties(shared.device).total_memory - torch.xpu.memory_allocated()), 'used': torch.xpu.memory_allocated(), 'total': torch.xpu.get_device_properties(shared.device).total_memory }
-                s = dict(torch.xpu.memory_stats())
-                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():
+            if 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))

modules/devices.py

@@ -8,7 +8,6 @@ from modules import cmd_args, shared, memstats
 if sys.platform == "darwin":
     from modules import mac_specific # pylint: disable=ungrouped-imports
 
-cuda_ok = torch.cuda.is_available()
 previous_oom = 0
 
 
@@ -192,10 +191,26 @@ else:
     backend = 'cpu'
 
 if backend == 'ipex':
-    #Fix broken functions with ipex
-    from modules.sd_hijack_utils import CondFunc
+    #Fix functions with ipex
+    torch.cuda.is_available = torch.xpu.is_available
+    torch.cuda.current_device = torch.xpu.current_device
+    torch.cuda.get_device_properties = torch.xpu.get_device_properties
     torch.cuda.empty_cache = torch_gc
 
+    torch.cuda.memory_stats = torch.xpu.memory_stats
+    torch.cuda.mem_get_info = lambda device: [(torch.xpu.get_device_properties(device).total_memory - torch.xpu.memory_allocated(device)), torch.xpu.get_device_properties(device).total_memory]
+    torch.cuda.memory_allocated = torch.xpu.memory_allocated
+    torch.cuda.max_memory_allocated = torch.xpu.max_memory_allocated
+    torch.cuda.reset_peak_memory_stats = torch.xpu.reset_peak_memory_stats
+
+    torch.cuda.get_rng_state_all = torch.xpu.get_rng_state_all
+    torch.cuda.set_rng_state_all = torch.xpu.set_rng_state_all
+    try:
+        torch.cuda.amp.GradScaler = torch.xpu.amp.GradScaler
+    except Exception:
+        pass
+
+    from modules.sd_hijack_utils import CondFunc
     #Functions with dtype errors:
     CondFunc('torch.nn.modules.GroupNorm.forward',
         lambda orig_func, *args, **kwargs: orig_func(args[0], args[1].to(args[0].weight.data.dtype)),
@@ -207,6 +222,7 @@ if backend == 'ipex':
     CondFunc('torch.nn.modules.Conv2d._conv_forward',
         lambda orig_func, *args, **kwargs: orig_func(args[0], args[1].to(args[2].data.dtype), args[2], args[3]),
         lambda *args, **kwargs: args[2].dtype != args[3].data.dtype)
+
     #Functions that does not work with the XPU:
     #UniPC:
     CondFunc('torch.linalg.solve',
@@ -238,6 +254,7 @@ if backend == 'ipex':
         args[5], args[6], args[7], args[8]).to(get_cuda_device_string()),
         lambda *args, **kwargs: args[1].device != torch.device("cpu"))
 
+cuda_ok = torch.cuda.is_available() and not backend == 'ipex'
 cpu = torch.device("cpu")
 device = device_interrogate = device_gfpgan = device_esrgan = device_codeformer = None
 dtype = torch.float16

modules/hypernetworks/hypernetwork.py

@@ -589,10 +589,7 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi
             print("Cannot resume from saved optimizer!")
             print(e)
 
-    if devices.backend == 'ipex':
-        scaler = torch.xpu.amp.GradScaler()
-    else:
-        scaler = torch.cuda.amp.GradScaler()
+    scaler = torch.cuda.amp.GradScaler()
 
     batch_size = ds.batch_size
     gradient_step = ds.gradient_step
@@ -706,10 +703,7 @@ 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 devices.backend == 'ipex':
-                        cuda_rng_state = torch.xpu.get_rng_state_all()
-                    elif torch.cuda.is_available():
-                        cuda_rng_state = torch.cuda.get_rng_state_all()
+                    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,10 +739,7 @@ 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 devices.backend == 'ipex':
-                        torch.xpu.set_rng_state_all(cuda_rng_state)
-                    elif torch.cuda.is_available():
-                        torch.cuda.set_rng_state_all(cuda_rng_state)
+                    torch.cuda.set_rng_state_all(cuda_rng_state)
                     hypernetwork.train()
                     if image is not None:
                         shared.state.assign_current_image(image)

modules/memmon.py

@@ -21,12 +21,7 @@ class MemUsageMonitor(threading.Thread):
         self.run_flag = threading.Event()
         self.data = defaultdict(int)
         if not torch.cuda.is_available():
-            #torch.cuda.is_available() reports False when using IPEX.
-            if devices.backend == 'ipex':
-                self.cuda_mem_get_info()
-                torch.xpu.memory_stats(self.device)
-            else:
-                self.disabled = True
+            self.disabled = True
         else:
             try:
                 self.cuda_mem_get_info()
@@ -35,22 +30,15 @@ class MemUsageMonitor(threading.Thread):
                 self.disabled = True
 
     def cuda_mem_get_info(self):
-        if devices.backend == 'ipex':
-            index = self.device.index if self.device.index is not None else torch.xpu.current_device()
-            return [(torch.xpu.get_device_properties(index).total_memory - torch.xpu.memory_allocated(index)), torch.xpu.get_device_properties(index).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)
+        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()
-            if devices.backend == 'ipex':
-                torch.xpu.reset_peak_memory_stats()
-            else:
-                torch.cuda.reset_peak_memory_stats()
+            torch.cuda.reset_peak_memory_stats()
             self.data.clear()
             if self.opts.memmon_poll_rate <= 0:
                 self.run_flag.clear()
@@ -70,10 +58,7 @@ class MemUsageMonitor(threading.Thread):
             self.data["free"] = free
             self.data["total"] = total
             try:
-                if devices.backend == 'ipex':
-                    torch_stats = torch.xpu.memory_stats(self.device)
-                else:
-                    torch_stats = torch.cuda.memory_stats(self.device)
+                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"]

modules/memstats.py

@@ -29,17 +29,4 @@ def memory_stats():
         return mem
     except Exception:
         pass
-    try:
-        gpu = { 'used': gb(torch.xpu.memory_allocated()), 'total': gb(torch.xpu.get_device_properties(shared.device).total_memory) }
-        s = dict(torch.xpu.memory_stats())
-        mem.update({
-            'gpu': gpu,
-            'retries': s['num_alloc_retries'],
-            'oom': s['num_ooms']
-        })
-        if s['num_ooms'] > 0:
-            shared.state.oom = True
-        return mem
-    except Exception:
-        pass
     return mem

modules/sd_hijack_optimizations.py

@@ -29,15 +29,7 @@ else:
 
 
 def get_available_vram():
-    if devices.backend == 'ipex':
-        stats = torch.xpu.memory_stats(shared.device)
-        mem_active = stats['active_bytes.all.current']
-        mem_reserved = stats['reserved_bytes.all.current']
-        mem_free_xpu = torch.xpu.get_device_properties(shared.device).total_memory - torch.xpu.memory_allocated(shared.device)
-        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':
+    if shared.device.type == 'cuda' or shared.device.type == 'xpu':
         try:
             stats = torch.cuda.memory_stats(shared.device)
             mem_active = stats['active_bytes.all.current']
@@ -47,7 +39,6 @@ def get_available_vram():
             mem_free_total = mem_free_cuda + mem_free_torch
         except Exception:
             mem_free_total = 1024 * 1024 * 1024
-
         return mem_free_total
     elif shared.device.type == 'privateuseone':
         mem_total, mem_active = torch.dml.memory_stats(shared.device)
@@ -190,27 +181,17 @@ 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):
-    if devices.backend == 'ipex':
-        stats = torch.xpu.memory_stats(q.device)
+    try:
+        stats = torch.cuda.memory_stats(q.device)
         mem_active = stats['active_bytes.all.current']
         mem_reserved = stats['reserved_bytes.all.current']
-        mem_free_xpu = torch.xpu.get_device_properties(q.device).total_memory - torch.xpu.memory_allocated(q.device)
+        mem_free_cuda, _ = torch.cuda.mem_get_info(q.device)
         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:
-        try:
-            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
-        except Exception:
-            mem_free_total = 1024 * 1024 * 1024
-        # 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))
+        mem_free_total = mem_free_cuda + mem_free_torch
+    except Exception:
+        mem_free_total = 1024 * 1024 * 1024
+    # 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)
@@ -218,7 +199,7 @@ 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 q.device.type == 'cuda' or devices.backend == 'ipex':
+    if q.device.type == 'cuda' or q.device.type == 'xpu':
         return einsum_op_cuda(q, k, v)
 
     if q.device.type == 'mps':
@@ -410,12 +391,8 @@ 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):
-    if devices.backend == '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)
+    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
@@ -523,12 +500,8 @@ def sdp_attnblock_forward(self, x):
     return x + out
 
 def sdp_no_mem_attnblock_forward(self, x):
-    if devices.backend == '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)
+    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

modules/sd_hijack_unet.py

@@ -67,7 +67,7 @@ def hijack_ddpm_edit():
 unet_needs_upcast = lambda *args, **kwargs: devices.unet_needs_upcast # pylint: disable=unnecessary-lambda-assignment
 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() or devices.backend == 'ipex':
+if version.parse(torch.__version__) <= version.parse("1.13.2") or torch.cuda.is_available():
     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)

modules/sd_models.py

@@ -737,7 +737,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
             sd_model.unet.to(memory_format=torch.channels_last)
 
         base_sent_to_cpu=False
-        if shared.opts.cuda_compile and (torch.cuda.is_available() or devices.backend == 'ipex'):
+        if shared.opts.cuda_compile and torch.cuda.is_available():
             if op == 'refiner':
                 gpu_vram = memory_stats().get('gpu', {})
                 free_vram = gpu_vram.get('total', 0) - gpu_vram.get('used', 0)

modules/textual_inversion/textual_inversion.py

@@ -443,10 +443,7 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st
         else:
             shared.log.info("No saved optimizer exists in checkpoint")
 
-    if devices.backend == 'ipex':
-        scaler = torch.xpu.amp.GradScaler()
-    else:
-        scaler = torch.cuda.amp.GradScaler()
+    scaler = torch.cuda.amp.GradScaler()
 
     batch_size = ds.batch_size
     gradient_step = ds.gradient_step