You can now select the device index with: --directml id

Like this for example: --directml 1
This commit is contained in:
comfyanonymous
2023-04-28 16:51:35 -04:00
parent cab80973d1
commit 2ca934f7d4
2 changed files with 10 additions and 4 deletions

View File

@@ -21,10 +21,15 @@ accelerate_enabled = False
xpu_available = False
directml_enabled = False
if args.directml:
if args.directml is not None:
import torch_directml
print("Using directml")
directml_enabled = True
device_index = args.directml
if device_index < 0:
directml_device = torch_directml.device()
else:
directml_device = torch_directml.device(device_index)
print("Using directml with device:", torch_directml.device_name(device_index))
# torch_directml.disable_tiled_resources(True)
try:
@@ -226,7 +231,8 @@ def get_torch_device():
global xpu_available
global directml_enabled
if directml_enabled:
return torch_directml.device()
global directml_device
return directml_device
if vram_state == VRAMState.MPS:
return torch.device("mps")
if vram_state == VRAMState.CPU: