diff --git a/utils.py b/utils.py index f8f9fe7..3bb65f8 100644 --- a/utils.py +++ b/utils.py @@ -8,7 +8,7 @@ from scipy.spatial import ConvexHull from folder_paths import models_dir from .BiSeNet import BiSeNet from ultralytics import YOLO -from onnxruntime import InferenceSession, get_available_providers +from onnxruntime import InferenceSession from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation from skimage import transform as trans @@ -62,7 +62,7 @@ class Models: @classmethod def lmk(cls, crop): if '_lmk' not in cls.__dict__: - cls._lmk = InferenceSession(os.path.join(models_dir, 'landmarks', 'fan2_68_landmark.onnx')) + cls._lmk = InferenceSession(os.path.join(models_dir, 'landmarks', 'fan2_68_landmark.onnx'), providers=['CUDAExecutionProvider', 'CPUExecutionProvider']) lmk = cls._lmk.run(None, {'input': crop})[0] return lmk