diff --git a/README.md b/README.md index ede89ec..2a7fb6a 100644 --- a/README.md +++ b/README.md @@ -5,6 +5,7 @@ These custom nodes provide a rotation aware face extraction, paste back, and var ![Comparison](examples/comparison.jpg) ## Patch notes +- 2024-06-14 - Sent yolo to cuda device if available (credits: kilroythethird, thanks for the suggestion!) - 2024-06-14 - Dropped the requirement from insightface completely. - 2024-05-22 - Updated GenderFaceFilter node. - 2024-05-19 - Added BiSeNetMask and JonathandinuMask nodes. Careful about JonathandinuMask, it's more accurate than BiSeNet, but it takes more memory; you can get out of memory more easily with it. diff --git a/utils.py b/utils.py index 679b5d1..35dccce 100644 --- a/utils.py +++ b/utils.py @@ -56,6 +56,8 @@ class Models: def yolo(cls, img, threshold): if '_yolo' not in cls.__dict__: cls._yolo = YOLO(os.path.join(models_dir,'ultralytics','bbox','face_yolov8m.pt')) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + cls._yolo = cls._yolo.to(device) dets = cls._yolo(img, conf=threshold)[0] return dets @classmethod