add Query Device

This commit is contained in:
ethan 2025-01-30 00:20:58 -08:00
parent 317af7201f
commit 77e9294c08

View File

@ -1,16 +1,25 @@
import torch import torch
import importlib
class TorchCompileModel: class TorchCompileModel:
@classmethod @classmethod
def INPUT_TYPES(s): def INPUT_TYPES(s):
if importlib.util.find_spec("openvino") is not None:
import openvino as ov
core = ov.Core()
available_devices = core.available_devices
else:
available_devices = []
return { return {
"required": { "required": {
"model": ("MODEL",), "model": ("MODEL",),
"backend": (["inductor", "cudagraphs", "openvino"],), "backend": (["inductor", "cudagraphs", "openvino"],),
}, },
"optional": { "optional": {
"openvino_device": (["CPU", "GPU", "NPU"],), "openvino_device": (available_devices,),
}, },
} }