93 lines
2.9 KiB
Python
93 lines
2.9 KiB
Python
# Dreamfast LTX2 Multi-GPU - Device Utilities
|
|
# Copyright (C) 2025 Dreamfast
|
|
# Based on ComfyUI-MultiGPU by pollockjj (GPL-3.0)
|
|
#
|
|
# This program is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU General Public License as published by
|
|
# the Free Software Foundation, either version 3 of the License, or
|
|
# (at your option) any later version.
|
|
|
|
import torch
|
|
import logging
|
|
import gc
|
|
|
|
logger = logging.getLogger("LTXV2MultiGPU")
|
|
|
|
_DEVICE_LIST_CACHE = None
|
|
|
|
|
|
def get_device_list():
|
|
"""
|
|
Enumerate all available devices that can store torch tensors.
|
|
Results are cached after first call since devices don't change during runtime.
|
|
|
|
Returns list like: ["cpu", "cuda:0", "cuda:1", "mps", ...]
|
|
"""
|
|
global _DEVICE_LIST_CACHE
|
|
|
|
if _DEVICE_LIST_CACHE is not None:
|
|
return _DEVICE_LIST_CACHE
|
|
|
|
devs = ["cpu"]
|
|
|
|
# CUDA (NVIDIA + AMD ROCm)
|
|
if hasattr(torch, "cuda") and torch.cuda.is_available():
|
|
device_count = torch.cuda.device_count()
|
|
devs += [f"cuda:{i}" for i in range(device_count)]
|
|
logger.debug(f"Found {device_count} CUDA device(s)")
|
|
|
|
# XPU (Intel)
|
|
try:
|
|
import intel_extension_for_pytorch
|
|
except ImportError:
|
|
pass
|
|
|
|
if hasattr(torch, "xpu") and hasattr(torch.xpu, "is_available") and torch.xpu.is_available():
|
|
device_count = torch.xpu.device_count()
|
|
devs += [f"xpu:{i}" for i in range(device_count)]
|
|
logger.debug(f"Found {device_count} XPU device(s)")
|
|
|
|
# MPS (Apple Metal)
|
|
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
|
devs.append("mps")
|
|
logger.debug("Found MPS device")
|
|
|
|
# DirectML (Windows)
|
|
try:
|
|
import torch_directml
|
|
adapter_count = torch_directml.device_count()
|
|
if adapter_count > 0:
|
|
devs += [f"directml:{i}" for i in range(adapter_count)]
|
|
logger.debug(f"Found {adapter_count} DirectML adapter(s)")
|
|
except ImportError:
|
|
pass
|
|
|
|
_DEVICE_LIST_CACHE = devs
|
|
logger.debug(f"Device list: {devs}")
|
|
|
|
return devs
|
|
|
|
|
|
def soft_empty_cache_all_devices():
|
|
"""Clear allocator caches across all devices."""
|
|
gc.collect()
|
|
|
|
all_devices = get_device_list()
|
|
is_cuda_available = hasattr(torch, "cuda") and torch.cuda.is_available()
|
|
|
|
for device_str in all_devices:
|
|
if device_str.startswith("cuda:") and is_cuda_available:
|
|
device_idx = int(device_str.split(":")[1])
|
|
with torch.cuda.device(device_idx):
|
|
torch.cuda.empty_cache()
|
|
if hasattr(torch.cuda, "ipc_collect"):
|
|
torch.cuda.ipc_collect()
|
|
|
|
elif device_str == "mps":
|
|
if hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"):
|
|
torch.mps.empty_cache()
|
|
|
|
elif device_str.startswith("xpu:"):
|
|
if hasattr(torch, "xpu") and hasattr(torch.xpu, "empty_cache"):
|
|
torch.xpu.empty_cache()
|