Update speed.py

This commit is contained in:
ZhizhouZhong 2024-07-13 14:45:33 +08:00 committed by GitHub
parent a96532752d
commit 286d1fd35a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -6,10 +6,13 @@ Benchmark the inference speed of each module in LivePortrait.
TODO: heavy GPT style, need to refactor TODO: heavy GPT style, need to refactor
""" """
import yaml
import torch import torch
torch._dynamo.config.suppress_errors = True # Suppress errors and fall back to eager execution
import yaml
import time import time
import numpy as np import numpy as np
from src.utils.helper import load_model, concat_feat from src.utils.helper import load_model, concat_feat
from src.config.inference_config import InferenceConfig from src.config.inference_config import InferenceConfig
@ -47,9 +50,6 @@ def load_and_compile_models(cfg, model_config):
""" """
Load and compile models for inference Load and compile models for inference
""" """
import torch._dynamo
torch._dynamo.config.suppress_errors = True # Suppress errors and fall back to eager execution
appearance_feature_extractor = load_model(cfg.checkpoint_F, model_config, cfg.device_id, 'appearance_feature_extractor') 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') 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') warping_module = load_model(cfg.checkpoint_W, model_config, cfg.device_id, 'warping_module')