OpenSourceJoker/LivePortrait
3
0
Fork 0
mirror of https://github.com/KwaiVGI/LivePortrait.git synced 2024-12-22 20:42:38 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix: compile on v100

Browse Source 
LGTM.
This commit is contained in:
ZhizhouZhong 2024-07-13 14:51:32 +08:00 committed by GitHub
commit 6bcbb014fb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 9 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
speed.py
View File

@ -6,10 +6,13 @@ Benchmark the inference speed of each module in LivePortrait.
TODO: heavy GPT style, need to refactor
"""
import yaml
import torch
torch._dynamo.config.suppress_errors = True # Suppress errors and fall back to eager execution
import yaml
import time
import numpy as np
from src.utils.helper import load_model, concat_feat
from src.config.inference_config import InferenceConfig
@ -47,11 +50,11 @@ def load_and_compile_models(cfg, model_config):
"""
Load and compile models for inference
"""
appearance_feature_extractor = load_model(cfg.checkpoint_F, model_config, cfg.device, 'appearance_feature_extractor')
motion_extractor = load_model(cfg.checkpoint_M, model_config, cfg.device, 'motion_extractor')
warping_module = load_model(cfg.checkpoint_W, model_config, cfg.device, 'warping_module')
spade_generator = load_model(cfg.checkpoint_G, model_config, cfg.device, 'spade_generator')
stitching_retargeting_module = load_model(cfg.checkpoint_S, model_config, cfg.device, 'stitching_retargeting_module')
appearance_feature_extractor = load_model(cfg.checkpoint_F, model_config, cfg.device_id, 'appearance_feature_extractor')
motion_extractor = load_model(cfg.checkpoint_M, model_config, cfg.device_id, 'motion_extractor')
warping_module = load_model(cfg.checkpoint_W, model_config, cfg.device_id, 'warping_module')
spade_generator = load_model(cfg.checkpoint_G, model_config, cfg.device_id, 'spade_generator')
stitching_retargeting_module = load_model(cfg.checkpoint_S, model_config, cfg.device_id, 'stitching_retargeting_module')
models_with_params = [
('Appearance Feature Extractor', appearance_feature_extractor),