nodes/liusida-ComfyUI-AutoCropFaces
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

use comfyui's device #5

Browse Source
This commit is contained in:
Sida Liu
2024-07-24 17:26:45 +08:00
parent 3fea2c662f
commit a62cdcb91a
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+2 -1
View File
@@ -1,6 +1,7 @@
import torch
import comfy.utils
from .Pytorch_Retinaface.pytorch_retinaface import Pytorch_RetinaFace
from comfy.model_management import get_torch_device
class AutoCropFaces:
def __init__(self):
@@ -67,7 +68,7 @@ class AutoCropFaces:
def auto_crop_faces_in_image (self, image, max_number_of_faces, scale_factor, shift_factor, aspect_ratio, method='lanczos'):
image_255 = image * 255
rf = Pytorch_RetinaFace(top_k=50, keep_top_k=max_number_of_faces)
rf = Pytorch_RetinaFace(top_k=50, keep_top_k=max_number_of_faces, device=get_torch_device())
dets = rf.detect_faces(image_255)
cropped_faces, bbox_info = rf.center_and_crop_rescale(image, dets, scale_factor=scale_factor, shift_factor=shift_factor, aspect_ratio=aspect_ratio)