improved Intel XPU compatibility

backend/memory_management.py

@@ -111,7 +111,10 @@ if args.directml is not None:
 
 try:
     import intel_extension_for_pytorch as ipex  # noqa: F401
+except Exception:
+    ipex = None
 
+try:
     _ = torch.xpu.device_count()
     xpu_available = torch.xpu.is_available()
 except Exception:
@@ -473,7 +476,7 @@ class LoadedModel:
 
         real_model = self.model.model
 
-        if is_intel_xpu() and not args.disable_ipex_optimize and "ipex" in globals() and real_model is not None:
+        if is_intel_xpu() and not args.disable_ipex_optimize and ipex is not None and real_model is not None:
             with torch.no_grad():
                 real_model = ipex.optimize(real_model.eval(), inplace=True, graph_mode=True, concat_linear=True)
 

backend/patcher/base.py

@@ -623,6 +623,8 @@ class ModelPatcher:
                 self.patch_weight_to_device(key, device_to=device_to)
             if memory_management.is_device_cuda(device_to):
                 torch.cuda.synchronize()
+            elif memory_management.is_device_xpu(device_to):
+                torch.xpu.synchronize()
 
             logger.debug("lowvram: loaded module regularly {} {}".format(n, m))
             m.forge_patched_weights = True

backend/quant_ops.py

@@ -21,10 +21,9 @@ else:
     if cuda_version < (13,):
         ck.registry.disable("cuda")
 
-try:
-    import triton  # noqa
-except Exception:
-    ck.registry.disable("triton")
+
+# https://github.com/Comfy-Org/ComfyUI/blob/v0.16.4/comfy/quant_ops.py#L24
+ck.registry.disable("triton")
 
 
 # region FP8 Layouts

backend/stream.py

@@ -19,6 +19,13 @@ def get_new_stream():
     return memory_management.get_offload_stream(device)
 
 
+def get_vae_stream():
+    if torch.cuda.is_available():
+        return torch.cuda.Stream(device=device, priority=1)
+    if torch.xpu.is_available():
+        return torch.xpu.Stream(device=device, priority=1)
+
+
 def should_use_stream():
     return current_stream is not None and mover_stream is not None
 

modules/api/api.py

@@ -766,15 +766,23 @@ class Api:
             import torch
 
             if torch.cuda.is_available():
-                s = torch.cuda.mem_get_info()
+                _backend = torch.cuda
+            elif torch.xpu.is_available():
+                _backend = torch.xpu
+            else:
+                _backend = None
+
+            if _backend is not None:
+                s = _backend.mem_get_info()
                 system = {"free": s[0], "used": s[1] - s[0], "total": s[1]}
-                s = dict(torch.cuda.memory_stats(shared.device))
+                s = dict(_backend.memory_stats(shared.device))
                 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"]}
+                active = {"current": s.get("active_bytes.all.current") or s.get("active.all.current") or 0, "peak": s.get("active_bytes.all.peak") or s.get("active.all.peak") or 0}
                 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 = {
+                    "device": str(shared.device),
                     "system": system,
                     "active": active,
                     "allocated": allocated,

modules/api/models.py

@@ -291,7 +291,7 @@ class EmbeddingsResponse(BaseModel):
 
 class MemoryResponse(BaseModel):
     ram: dict = Field(title="RAM", description="System memory stats")
-    cuda: dict = Field(title="CUDA", description="nVidia CUDA memory stats")
+    cuda: dict = Field(title="VRAM", description="GPU memory stats")
 
 
 class ScriptsList(BaseModel):

modules/launch_utils.py

@@ -319,11 +319,18 @@ def prepare_environment():
         startup_timer.record("install torch")
 
     if not args.skip_torch_cuda_test:
-        success, err = check_run_python("import torch; assert torch.cuda.is_available()", return_error=True)
+        TORCH_CHECK: str = """
+import torch
+cuda = hasattr(torch, "cuda") and torch.cuda.is_available()
+xpu = hasattr(torch, "xpu") and torch.xpu.is_available()
+assert cuda or xpu
+        """
+
+        success, err = check_run_python(TORCH_CHECK, return_error=True)
         if not success:
             if "older driver" in str(err).lower():
                 raise SystemError("Please update your GPU driver to support cu130 ; or manually install older PyTorch")
-            raise RuntimeError("PyTorch is not able to access CUDA")
+            raise RuntimeError("PyTorch is not able to access GPU")
         startup_timer.record("torch GPU test")
 
     if not is_installed("packaging"):

modules/memmon.py

@@ -4,6 +4,8 @@ from collections import defaultdict
 
 import torch
 
+from backend import memory_management
+
 
 class MemUsageMonitor(threading.Thread):
     run_flag = None
@@ -22,16 +24,21 @@ class MemUsageMonitor(threading.Thread):
         self.run_flag = threading.Event()
         self.data = defaultdict(int)
 
+        if memory_management.is_intel_xpu():
+            self._backend = torch.xpu
+        else:
+            self._backend = torch.cuda
+
         try:
             self.cuda_mem_get_info()
-            torch.cuda.memory_stats(self.device)
+            self._backend.memory_stats(self.device)
         except Exception as e:  # AMD or whatever
-            print(f"Warning: caught exception '{e}', memory monitor disabled")
+            memory_management.logger.warning(f'Caught Exception "{e}"\nMemory Monitor Disabled...')
             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)
+        index = self.device.index if self.device.index is not None else self._backend.current_device()
+        return self._backend.mem_get_info(index)
 
     def run(self):
         if self.disabled:
@@ -40,7 +47,7 @@ class MemUsageMonitor(threading.Thread):
         while True:
             self.run_flag.wait()
 
-            torch.cuda.reset_peak_memory_stats()
+            self._backend.reset_peak_memory_stats()
             self.data.clear()
 
             if self.opts.memmon_poll_rate <= 0:
@@ -61,13 +68,13 @@ class MemUsageMonitor(threading.Thread):
             print(k, -(v // -(1024**2)))
 
         print(self, "raw torch memory stats:")
-        tm = torch.cuda.memory_stats(self.device)
+        tm = self._backend.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())
+        print(self._backend.memory_summary())
 
     def monitor(self):
         self.run_flag.set()
@@ -78,9 +85,9 @@ class MemUsageMonitor(threading.Thread):
             self.data["free"] = free
             self.data["total"] = total
 
-            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"]
+            torch_stats = self._backend.memory_stats(self.device)
+            self.data["active"] = torch_stats.get("active_bytes.all.current") or torch_stats.get("active.all.current") or 0
+            self.data["active_peak"] = torch_stats.get("active_bytes.all.peak") or torch_stats.get("active.all.peak") or 0
             self.data["reserved"] = torch_stats["reserved_bytes.all.current"]
             self.data["reserved_peak"] = torch_stats["reserved_bytes.all.peak"]
             self.data["system_peak"] = total - self.data["min_free"]

modules/shared_state.py

@@ -38,9 +38,10 @@ class State:
 
     def __init__(self):
         self.server_start = time.time()
-        self.vae_stream = None
         if stream.should_use_stream():
-            self.vae_stream = torch.cuda.Stream(device=devices.device, priority=1)
+            self.vae_stream = stream.get_vae_stream()
+        else:
+            self.vae_stream = None
 
     @property
     def need_restart(self) -> bool: