29 lines
897 B
Python
29 lines
897 B
Python
import comfy
|
|
import comfy.model_management
|
|
|
|
# Patched version of comfy.model_management.get_torch_device that allows vulkan devices to be treated as default
|
|
def get_torch_device_patched():
|
|
global directml_enabled
|
|
global cpu_state
|
|
if directml_enabled:
|
|
global directml_device
|
|
return directml_device
|
|
if cpu_state == CPUState.MPS:
|
|
return torch.device("mps")
|
|
if cpu_state == CPUState.CPU:
|
|
return torch.device("cpu")
|
|
else:
|
|
if True:
|
|
# for now, just see if this works c:
|
|
return torch.device(f"vulkan:0")
|
|
else:
|
|
if is_intel_xpu():
|
|
return torch.device("xpu", torch.xpu.current_device())
|
|
elif is_ascend_npu():
|
|
return torch.device("npu", torch.npu.current_device())
|
|
elif is_mlu():
|
|
return torch.device("mlu", torch.mlu.current_device())
|
|
else:
|
|
return torch.device(torch.cuda.current_device())
|
|
comfy.model_management.get_torch_device = get_torch_device_patched
|