fix: torch.backends check (#280)

This commit is contained in:
ZhizhouZhong 2024-08-05 14:17:17 +08:00 committed by GitHub
parent 67d567f38c
commit 357226b2e7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 25 additions and 16 deletions

View File

@ -32,10 +32,13 @@ class LivePortraitWrapper(object):
if inference_cfg.flag_force_cpu:
self.device = 'cpu'
else:
try:
if torch.backends.mps.is_available():
self.device = 'mps'
else:
self.device = 'cuda:' + str(self.device_id)
except:
self.device = 'cuda:' + str(self.device_id)
model_config = yaml.load(open(inference_cfg.models_config, 'r'), Loader=yaml.SafeLoader)
# init F
@ -344,10 +347,13 @@ class LivePortraitWrapperAnimal(LivePortraitWrapper):
if inference_cfg.flag_force_cpu:
self.device = 'cpu'
else:
try:
if torch.backends.mps.is_available():
self.device = 'mps'
else:
self.device = 'cuda:' + str(self.device_id)
except:
self.device = 'cuda:' + str(self.device_id)
model_config = yaml.load(open(inference_cfg.models_config, 'r'), Loader=yaml.SafeLoader)
# init F

View File

@ -48,6 +48,7 @@ class Cropper(object):
device = "cpu"
face_analysis_wrapper_provider = ["CPUExecutionProvider"]
else:
try:
if torch.backends.mps.is_available():
# Shape inference currently fails with CoreMLExecutionProvider
# for the retinaface model
@ -56,7 +57,9 @@ class Cropper(object):
else:
device = "cuda"
face_analysis_wrapper_provider = ["CUDAExecutionProvider"]
except:
device = "cuda"
face_analysis_wrapper_provider = ["CUDAExecutionProvider"]
self.face_analysis_wrapper = FaceAnalysisDIY(
name="buffalo_l",
root=self.crop_cfg.insightface_root,