From d099ecd4971ee51dc53240b2a1e725537860c2b9 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Fri, 27 Dec 2024 03:06:50 -0600 Subject: [PATCH] Enhance device selection in get_torch_device_patched and override_class to include 'cpu' as a valid option for multi-GPU support --- __init__.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/__init__.py b/__init__.py index c27e385..8257316 100644 --- a/__init__.py +++ b/__init__.py @@ -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"