KlingTeam/LivePortrait

1# coding: utf-82 3"""4Benchmark the inference speed of each module in LivePortrait.5 6TODO: heavy GPT style, need to refactor7"""8 9import yaml10import torch11import time12import numpy as np13from src.utils.helper import load_model, concat_feat14from src.config.inference_config import InferenceConfig15 16 17def initialize_inputs(batch_size=1):18    """19    Generate random input tensors and move them to GPU20    """21    feature_3d = torch.randn(batch_size, 32, 16, 64, 64).cuda().half()22    kp_source = torch.randn(batch_size, 21, 3).cuda().half()23    kp_driving = torch.randn(batch_size, 21, 3).cuda().half()24    source_image = torch.randn(batch_size, 3, 256, 256).cuda().half()25    generator_input = torch.randn(batch_size, 256, 64, 64).cuda().half()26    eye_close_ratio = torch.randn(batch_size, 3).cuda().half()27    lip_close_ratio = torch.randn(batch_size, 2).cuda().half()28    feat_stitching = concat_feat(kp_source, kp_driving).half()29    feat_eye = concat_feat(kp_source, eye_close_ratio).half()30    feat_lip = concat_feat(kp_source, lip_close_ratio).half()31 32    inputs = {33        'feature_3d': feature_3d,34        'kp_source': kp_source,35        'kp_driving': kp_driving,36        'source_image': source_image,37        'generator_input': generator_input,38        'feat_stitching': feat_stitching,39        'feat_eye': feat_eye,40        'feat_lip': feat_lip41    }42 43    return inputs44 45 46def load_and_compile_models(cfg, model_config):47    """48    Load and compile models for inference49    """50    appearance_feature_extractor = load_model(cfg.checkpoint_F, model_config, cfg.device_id, 'appearance_feature_extractor')51    motion_extractor = load_model(cfg.checkpoint_M, model_config, cfg.device_id, 'motion_extractor')52    warping_module = load_model(cfg.checkpoint_W, model_config, cfg.device_id, 'warping_module')53    spade_generator = load_model(cfg.checkpoint_G, model_config, cfg.device_id, 'spade_generator')54    stitching_retargeting_module = load_model(cfg.checkpoint_S, model_config, cfg.device_id, 'stitching_retargeting_module')55 56    models_with_params = [57        ('Appearance Feature Extractor', appearance_feature_extractor),58        ('Motion Extractor', motion_extractor),59        ('Warping Network', warping_module),60        ('SPADE Decoder', spade_generator)61    ]62 63    compiled_models = {}64    for name, model in models_with_params:65        model = model.half()66        model = torch.compile(model, mode='max-autotune')  # Optimize for inference67        model.eval()  # Switch to evaluation mode68        compiled_models[name] = model69 70    retargeting_models = ['stitching', 'eye', 'lip']71    for retarget in retargeting_models:72        module = stitching_retargeting_module[retarget].half()73        module = torch.compile(module, mode='max-autotune')  # Optimize for inference74        module.eval()  # Switch to evaluation mode75        stitching_retargeting_module[retarget] = module76 77    return compiled_models, stitching_retargeting_module78 79 80def warm_up_models(compiled_models, stitching_retargeting_module, inputs):81    """82    Warm up models to prepare them for benchmarking83    """84    print("Warm up start!")85    with torch.no_grad():86        for _ in range(10):87            compiled_models['Appearance Feature Extractor'](inputs['source_image'])88            compiled_models['Motion Extractor'](inputs['source_image'])89            compiled_models['Warping Network'](inputs['feature_3d'], inputs['kp_driving'], inputs['kp_source'])90            compiled_models['SPADE Decoder'](inputs['generator_input'])  # Adjust input as required91            stitching_retargeting_module['stitching'](inputs['feat_stitching'])92            stitching_retargeting_module['eye'](inputs['feat_eye'])93            stitching_retargeting_module['lip'](inputs['feat_lip'])94    print("Warm up end!")95 96 97def measure_inference_times(compiled_models, stitching_retargeting_module, inputs):98    """99    Measure inference times for each model100    """101    times = {name: [] for name in compiled_models.keys()}102    times['Retargeting Models'] = []103 104    overall_times = []105 106    with torch.no_grad():107        for _ in range(100):108            torch.cuda.synchronize()109            overall_start = time.time()110 111            start = time.time()112            compiled_models['Appearance Feature Extractor'](inputs['source_image'])113            torch.cuda.synchronize()114            times['Appearance Feature Extractor'].append(time.time() - start)115 116            start = time.time()117            compiled_models['Motion Extractor'](inputs['source_image'])118            torch.cuda.synchronize()119            times['Motion Extractor'].append(time.time() - start)120 121            start = time.time()122            compiled_models['Warping Network'](inputs['feature_3d'], inputs['kp_driving'], inputs['kp_source'])123            torch.cuda.synchronize()124            times['Warping Network'].append(time.time() - start)125 126            start = time.time()127            compiled_models['SPADE Decoder'](inputs['generator_input'])  # Adjust input as required128            torch.cuda.synchronize()129            times['SPADE Decoder'].append(time.time() - start)130 131            start = time.time()132            stitching_retargeting_module['stitching'](inputs['feat_stitching'])133            stitching_retargeting_module['eye'](inputs['feat_eye'])134            stitching_retargeting_module['lip'](inputs['feat_lip'])135            torch.cuda.synchronize()136            times['Retargeting Models'].append(time.time() - start)137 138            overall_times.append(time.time() - overall_start)139 140    return times, overall_times141 142 143def print_benchmark_results(compiled_models, stitching_retargeting_module, retargeting_models, times, overall_times):144    """145    Print benchmark results with average and standard deviation of inference times146    """147    average_times = {name: np.mean(times[name]) * 1000 for name in times.keys()}148    std_times = {name: np.std(times[name]) * 1000 for name in times.keys()}149 150    for name, model in compiled_models.items():151        num_params = sum(p.numel() for p in model.parameters())152        num_params_in_millions = num_params / 1e6153        print(f"Number of parameters for {name}: {num_params_in_millions:.2f} M")154 155    for index, retarget in enumerate(retargeting_models):156        num_params = sum(p.numel() for p in stitching_retargeting_module[retarget].parameters())157        num_params_in_millions = num_params / 1e6158        print(f"Number of parameters for part_{index} in Stitching and Retargeting Modules: {num_params_in_millions:.2f} M")159 160    for name, avg_time in average_times.items():161        std_time = std_times[name]162        print(f"Average inference time for {name} over 100 runs: {avg_time:.2f} ms (std: {std_time:.2f} ms)")163 164 165def main():166    """167    Main function to benchmark speed and model parameters168    """169    # Sample input tensors170    inputs = initialize_inputs()171 172    # Load configuration173    cfg = InferenceConfig(device_id=0)174    model_config_path = cfg.models_config175    with open(model_config_path, 'r') as file:176        model_config = yaml.safe_load(file)177 178    # Load and compile models179    compiled_models, stitching_retargeting_module = load_and_compile_models(cfg, model_config)180 181    # Warm up models182    warm_up_models(compiled_models, stitching_retargeting_module, inputs)183 184    # Measure inference times185    times, overall_times = measure_inference_times(compiled_models, stitching_retargeting_module, inputs)186 187    # Print benchmark results188    print_benchmark_results(compiled_models, stitching_retargeting_module, ['stitching', 'eye', 'lip'], times, overall_times)189 190 191if __name__ == "__main__":192    main()193