onnx device select
This commit is contained in:
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user