xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-02 20:37:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Added support for Moore Threads GPUs

Browse Source
This commit is contained in:
hben35096 2025-05-09 02:26:43 +08:00 committed by GitHub
parent e56e423e44
commit d22feacb05
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 37 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
comfy/model_management.py
View File

@ -128,6 +128,15 @@ try:
except: except:
mlu_available = False mlu_available = False
try:
import torch_musa
_ = torch_musa.device_count()
musa_available = torch_musa.is_available()
if musa_available:
logging.info("MUSA device detected: {}".format(torch_musa.get_device_name(0)))
except:
musa_available = False
if args.cpu: if args.cpu:
cpu_state = CPUState.CPU cpu_state = CPUState.CPU
@ -151,6 +160,12 @@ def is_mlu():
return True return True
return False return False
def is_musa():
global musa_available
if musa_available:
return True
return False
def get_torch_device(): def get_torch_device():
global directml_enabled global directml_enabled
global cpu_state global cpu_state
@ -168,6 +183,8 @@ def get_torch_device():
return torch.device("npu", torch.npu.current_device()) return torch.device("npu", torch.npu.current_device())
elif is_mlu(): elif is_mlu():
return torch.device("mlu", torch.mlu.current_device()) return torch.device("mlu", torch.mlu.current_device())
elif is_musa():
return torch.device('musa', torch.musa.current_device())
else: else:
return torch.device(torch.cuda.current_device()) return torch.device(torch.cuda.current_device())
@ -200,6 +217,12 @@ def get_total_memory(dev=None, torch_total_too=False):
_, mem_total_mlu = torch.mlu.mem_get_info(dev) _, mem_total_mlu = torch.mlu.mem_get_info(dev)
mem_total_torch = mem_reserved mem_total_torch = mem_reserved
mem_total = mem_total_mlu mem_total = mem_total_mlu
elif is_musa():
stats = torch.musa.memory_stats(dev)
mem_reserved = stats['reserved_bytes.all.current']
_, mem_total = torch.musa.mem_get_info(dev)
mem_total_torch = mem_reserved
else: else:
stats = torch.cuda.memory_stats(dev) stats = torch.cuda.memory_stats(dev)
mem_reserved = stats['reserved_bytes.all.current'] mem_reserved = stats['reserved_bytes.all.current']
@ -1099,6 +1122,14 @@ def get_free_memory(dev=None, torch_free_too=False):
mem_free_mlu, _ = torch.mlu.mem_get_info(dev) mem_free_mlu, _ = torch.mlu.mem_get_info(dev)
mem_free_torch = mem_reserved - mem_active mem_free_torch = mem_reserved - mem_active
mem_free_total = mem_free_mlu + mem_free_torch mem_free_total = mem_free_mlu + mem_free_torch
elif is_musa():
stats = torch.musa.memory_stats(dev)
mem_active = stats['active_bytes.all.current']
mem_reserved = stats['reserved_bytes.all.current']
mem_free_musa, _ = torch.musa.mem_get_info(dev)
mem_free_torch = mem_reserved - mem_active
mem_free_total = mem_free_musa + mem_free_torch
else: else:
stats = torch.cuda.memory_stats(dev) stats = torch.cuda.memory_stats(dev)
mem_active = stats['active_bytes.all.current'] mem_active = stats['active_bytes.all.current']
@ -1171,6 +1202,9 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma
if is_mlu(): if is_mlu():
return True return True
if is_musa():
return True
if torch.version.hip: if torch.version.hip:
return True return True
@ -1231,6 +1265,9 @@ def should_use_bf16(device=None, model_params=0, prioritize_performance=True, ma
if is_ascend_npu(): if is_ascend_npu():
return True return True
if is_musa():
return True
if is_amd(): if is_amd():
arch = torch.cuda.get_device_properties(device).gcnArchName arch = torch.cuda.get_device_properties(device).gcnArchName
if any((a in arch) for a in ["gfx1030", "gfx1031", "gfx1010", "gfx1011", "gfx1012", "gfx906", "gfx900", "gfx803"]): # RDNA2 and older don't support bf16 if any((a in arch) for a in ["gfx1030", "gfx1031", "gfx1010", "gfx1011", "gfx1012", "gfx906", "gfx900", "gfx803"]): # RDNA2 and older don't support bf16