Enhance device selection in get_torch_device_patched and override_class to include 'cpu' as a valid option for multi-GPU support

This commit is contained in:
John Pollock
2024-12-27 03:06:50 -06:00
parent dead358227
commit d099ecd497
+3 -1
View File
@@ -9,6 +9,7 @@ def get_torch_device_patched():
if (
not torch.cuda.is_available()
or comfy.model_management.cpu_state == comfy.model_management.CPUState.CPU
or "cpu" in str(current_device).lower()
):
return torch.device("cpu")
return torch.device(current_device)
@@ -20,7 +21,8 @@ def override_class(cls):
@classmethod
def INPUT_TYPES(s):
inputs = copy.deepcopy(cls.INPUT_TYPES())
inputs["required"]["device"] = ([f"cuda:{i}" for i in range(torch.cuda.device_count())],)
devices = ["cpu"] + [f"cuda:{i}" for i in range(torch.cuda.device_count())]
inputs["required"]["device"] = (devices,)
return inputs
CATEGORY = "multigpu"