OSSForks/LivePortrait
1
0
Fork 0
mirror of https://github.com/zebrajr/LivePortrait.git synced 2026-01-15 12:15:12 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix: torch.backends check (#280)

Browse Source
This commit is contained in:
ZhizhouZhong
2024-08-05 14:17:17 +08:00
committed by GitHub
parent 67d567f38c
commit 357226b2e7
2 changed files with 25 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
src/live_portrait_wrapper.py
View File

@@ -32,9 +32,12 @@ class LivePortraitWrapper(object):
if inference_cfg.flag_force_cpu:
self.device = 'cpu'
else:
if torch.backends.mps.is_available():
self.device = 'mps'
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)
@@ -344,10 +347,13 @@ class LivePortraitWrapperAnimal(LivePortraitWrapper):
if inference_cfg.flag_force_cpu:
self.device = 'cpu'
else:
if torch.backends.mps.is_available():
self.device = 'mps'
else:
self.device = 'cuda:' + str(self.device_id)
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

21
src/utils/cropper.py
View File

@@ -48,15 +48,18 @@ class Cropper(object):
device = "cpu"
face_analysis_wrapper_provider = ["CPUExecutionProvider"]
else:
if torch.backends.mps.is_available():
# Shape inference currently fails with CoreMLExecutionProvider
# for the retinaface model
device = "mps"
face_analysis_wrapper_provider = ["CPUExecutionProvider"]
else:
device = "cuda"
face_analysis_wrapper_provider = ["CUDAExecutionProvider"]
try:
if torch.backends.mps.is_available():
# Shape inference currently fails with CoreMLExecutionProvider
# for the retinaface model
device = "mps"
face_analysis_wrapper_provider = ["CPUExecutionProvider"]
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,