diff --git a/backend/memory_management.py b/backend/memory_management.py index 2ceee434..ca06ff36 100644 --- a/backend/memory_management.py +++ b/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) diff --git a/backend/patcher/base.py b/backend/patcher/base.py index c0d5af29..2a264ee2 100644 --- a/backend/patcher/base.py +++ b/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 diff --git a/backend/quant_ops.py b/backend/quant_ops.py index 28a133f0..59a44fd5 100644 --- a/backend/quant_ops.py +++ b/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 diff --git a/backend/stream.py b/backend/stream.py index 8c85e678..0fd3e7f0 100644 --- a/backend/stream.py +++ b/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 diff --git a/modules/api/api.py b/modules/api/api.py index a454efe7..36ca1517 100644 --- a/modules/api/api.py +++ b/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, diff --git a/modules/api/models.py b/modules/api/models.py index 0a602450..16421b81 100644 --- a/modules/api/models.py +++ b/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): diff --git a/modules/launch_utils.py b/modules/launch_utils.py index 7ec2bae5..db54acf9 100644 --- a/modules/launch_utils.py +++ b/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"): diff --git a/modules/memmon.py b/modules/memmon.py index 4ce84c2a..3e15ae04 100644 --- a/modules/memmon.py +++ b/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"] diff --git a/modules/shared_state.py b/modules/shared_state.py index 513271b5..6d7db863 100644 --- a/modules/shared_state.py +++ b/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: