onnx device select

This commit is contained in:
kijai
2024-07-08 16:45:33 +03:00
parent 65574a6b59
commit 2533e6ad95
2 changed files with 16 additions and 5 deletions
+3 -3
View File
@@ -34,12 +34,12 @@ class Trajectory:
class Cropper(object):
def __init__(self, **kwargs) -> None:
def __init__(self, provider, **kwargs) -> None:
device_id = kwargs.get('device_id', 0)
self.landmark_runner = LandmarkRunner(
#ckpt_path=make_abs_path('../../pretrained_weights/liveportrait/landmark.onnx'),
ckpt_path=os.path.join(folder_paths.models_dir, 'liveportrait', 'landmark.onnx'),
onnx_provider='cuda',
onnx_provider=provider,
device_id=device_id
)
self.landmark_runner.warmup()
@@ -47,7 +47,7 @@ class Cropper(object):
self.face_analysis_wrapper = FaceAnalysisDIY(
name='buffalo_l',
root=os.path.join(folder_paths.models_dir, 'insightface'),
providers=["CUDAExecutionProvider"]
providers=[provider + 'ExecutionProvider',]
)
self.face_analysis_wrapper.prepare(ctx_id=device_id, det_size=(512, 512))
self.face_analysis_wrapper.warmup()
+13 -2
View File
@@ -237,6 +237,17 @@ class LivePortraitProcess:
"stitching": ("BOOLEAN", {"default": True}),
"relative": ("BOOLEAN", {"default": True}),
},
"optional": {
"onnx_device": (
[
'CUDA',
'CPU',
], {
"default": 'CUDA'
}),
}
}
RETURN_TYPES = ("IMAGE", "IMAGE",)
@@ -245,7 +256,7 @@ class LivePortraitProcess:
CATEGORY = "LivePortrait"
def process(self, source_image, driving_images, dsize, scale, vx_ratio, vy_ratio, pipeline,
lip_zero, eye_retargeting, lip_retargeting, stitching, relative, eyes_retargeting_multiplier, lip_retargeting_multiplier):
lip_zero, eye_retargeting, lip_retargeting, stitching, relative, eyes_retargeting_multiplier, lip_retargeting_multiplier, onnx_device='CUDA'):
source_image_np = (source_image * 255).byte().numpy()
driving_images_np = (driving_images * 255).byte().numpy()
@@ -256,7 +267,7 @@ class LivePortraitProcess:
vy_ratio = vy_ratio,
)
cropper = Cropper(crop_cfg=crop_cfg)
cropper = Cropper(crop_cfg=crop_cfg, provider=onnx_device)
pipeline.cropper = cropper
pipeline.live_portrait_wrapper.cfg.flag_eye_retargeting = eye_retargeting
pipeline.live_portrait_wrapper.cfg.eyes_retargeting_multiplier = eyes_retargeting_multiplier