Sent yolo to cuda device if available

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.

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