mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-04-30 20:32:45 +08:00
Update model_management.py
fallback to cpu if nvidia is not available
This commit is contained in:
parent
c033bbf516
commit
cbd3ba3dec
@ -203,8 +203,10 @@ 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())
|
||||||
else:
|
elif torch.cuda.is_available():
|
||||||
return torch.device(torch.cuda.current_device())
|
return torch.device(torch.cuda.current_device())
|
||||||
|
else:
|
||||||
|
return torch.device("cpu")
|
||||||
|
|
||||||
def get_total_memory(dev=None, torch_total_too=False):
|
def get_total_memory(dev=None, torch_total_too=False):
|
||||||
global directml_enabled
|
global directml_enabled
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user