search

dchatel / comfyui_facetools

These custom nodes provide a rotation aware face extraction, paste back, and various face related masking options.
MIT License
45 stars 7 forks source link

Yolo doesn't run on GPU with rocm #11

Closed kilroythethird closed 1 week ago

kilroythethird commented 1 week ago

Hey, first of all thanks for this node. It improved my results noticeable.

On my system the YOLO model runs on CPU by default even if there is a cuda backend available (via rocm). If possible, please add a `.to'' call after loading the model. This fixes the issue with no downsides as far as i see.

In Models.yolo:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
cls._yolo = cls._yolo.to(device)

This might be a theoretical issue for other torch backends but from what i see yolo defaults to CPU for these too, so at least this shouldn't make it worse than it was before.

dchatel commented 1 week ago

Done. Thanks. I don't why I haven't set that earlier.