wehub-resource-sync
73 lines
3.4 KiB
Python
73 lines
3.4 KiB
Python
import torch
|
|
import argparse
|
|
import pathlib
|
|
import re
|
|
|
|
|
|
def modify_spk_embed(spk_embed):
|
|
num_spk, hidden_size = spk_embed.shape
|
|
all_ids = set(range(num_spk))
|
|
if args.drop is not None:
|
|
drop_ids = set([int(i) for i in args.drop.split(',') if i != '']).intersection(all_ids)
|
|
else:
|
|
drop_ids = all_ids - set([int(i) for i in args.retain.split(',') if i != ''])
|
|
|
|
fill_list = None
|
|
if args.fill == 'zeros':
|
|
fill_list = [0. for _ in drop_ids]
|
|
elif args.fill == 'random':
|
|
fill_list = [torch.randn(1, hidden_size, dtype=torch.float32, device='cpu') for _ in drop_ids]
|
|
elif args.fill == 'mean':
|
|
mean = torch.mean(spk_embed, dim=0, keepdim=True)
|
|
fill_list = [mean for _ in drop_ids]
|
|
elif args.fill == 'cyclic':
|
|
retain_ids = sorted(all_ids - drop_ids)
|
|
num_retain = len(retain_ids)
|
|
fill_list = [spk_embed[retain_ids[i % num_retain], :] for i, _ in enumerate(drop_ids)]
|
|
|
|
for spk_id, fill in zip(sorted(drop_ids), fill_list):
|
|
spk_embed[spk_id, :] = fill
|
|
|
|
|
|
parser = argparse.ArgumentParser(description='Drop or edit spk_embed in a checkpoint.')
|
|
parser.add_argument('input', type=str, help='Path to the input file')
|
|
parser.add_argument('output', type=str, help='Path to the output file')
|
|
drop_retain_group = parser.add_mutually_exclusive_group()
|
|
drop_retain_group.add_argument('--drop', type=str, required=False, metavar='ID,ID,...',
|
|
help='Drop specific speaker IDs.')
|
|
drop_retain_group.add_argument('--retain', type=str, required=False, metavar='ID,ID,...',
|
|
help='Retain specific speaker IDs and drop all the others.')
|
|
parser.add_argument('--fill', type=str, required=False, default='zeros', metavar='METHOD',
|
|
choices=['zeros', 'random', 'mean', 'cyclic'],
|
|
help='Specify a filling method for the dropped embedding. '
|
|
'Available methods: zeros, random, mean, cyclic')
|
|
parser.add_argument('--overwrite', required=False, default=False,
|
|
action='store_true', help='Overwrite if the output file exists.')
|
|
args = parser.parse_args()
|
|
assert args.drop is not None or args.retain is not None, 'Either --drop or --retain should be specified.'
|
|
if args.drop and not re.fullmatch(r'(\d+)?(,\d+)*,?', args.drop):
|
|
print(f'Invalid format for --drop: \'{args.drop}\'')
|
|
exit(-1)
|
|
if args.retain and not re.fullmatch(r'(\d+)?(,\d+)*,?', args.retain):
|
|
print(f'Invalid format for --retain: \'{args.retain}\'')
|
|
exit(-1)
|
|
|
|
import torch
|
|
input_ckpt = pathlib.Path(args.input).resolve()
|
|
output_ckpt = pathlib.Path(args.output).resolve()
|
|
assert input_ckpt.exists(), 'The input file does not exist.'
|
|
assert args.overwrite or not output_ckpt.exists(), \
|
|
'The output file already exists or is the same as the input file.\n' \
|
|
'This is not recommended because spk_embed dropping scripts may not be stable, ' \
|
|
'and you may be at risk of losing your model.\n' \
|
|
'If you are sure to OVERWRITE the existing file, please re-run this script with the \'--overwrite\' argument.'
|
|
|
|
ckpt_loaded = torch.load(input_ckpt, map_location='cpu')
|
|
state_dict = ckpt_loaded['state_dict']
|
|
if 'model.fs2.spk_embed.weight' in state_dict:
|
|
modify_spk_embed(state_dict['model.fs2.spk_embed.weight'])
|
|
if 'model.spk_embed.weight' in state_dict:
|
|
modify_spk_embed(state_dict['model.spk_embed.weight'])
|
|
|
|
torch.save(ckpt_loaded, output_ckpt)
|