import torch
ckpt_path = './baseTTS/SP/G.pth'
ckpt = torch.load(ckpt_path, map_location='cpu')
print(ckpt.keys())

new_ckpt = {}
new_ckpt['model'] = ckpt['model']
torch.save(new_ckpt, f'{ckpt_path}_new')

# # for vc model, delete sdp and dp weight
# ckpt_path = './multilingual_ckpts/converter/checkpoint.pth'
# ckpt = torch.load(ckpt_path, map_location='cpu')
# new_ckpt = {}
# model = {}
# for key, value in ckpt['model'].items():
#     if 'dp' in key or 'sdp' in key or 'enc_p' in key:
#         continue
#     model[key] = value
# new_ckpt = {'model': model}
# torch.save(new_ckpt, f'{ckpt_path}')

# from api import ToneColorConverter
# tone_color_converter = ToneColorConverter(f'./multilingual_ckpts/converter/config.json', device='cpu')
# tone_color_converter.load_ckpt(f'./multilingual_ckpts/converter/checkpoint.pth')
