model_type = args.model or 'vqvae'
if args.name:
model_name = f'{model_type}.{args.name}'
else:
model_name = f'{model_type}'
if model_type[:5] == 'vqvae':
print("Model type is vqvae")
model_fn = lambda dataset: vqvae.Model(model_type=model_type,
rnn_dims=896,
fc_dims=896,
global_decoder_cond_dims=dataset.
num_speakers(),
upsample_factors=(4, 4, 4),
num_group=args.num_group,
num_sample=args.num_sample,
normalize_vq=True,
noise_x=True,
noise_y=True).cuda()
dataset_type = 'multi'
elif model_type == 'wavernn':
print("Model type is wavernn")
model_fn = lambda dataset: wr.Model(rnn_dims=896,
fc_dims=896,
pad=2,
upsample_factors=(4, 4, 4),
feat_dims=80).cuda()
dataset_type = 'single'
elif model_type == 'nc':
model_name = f'{model_type}.43.upconv'
data_path = config.multi_speaker_data_path
with open(f'{data_path}/index.pkl', 'rb') as f:
index = pickle.load(f)
test_index = [
x[-1:] if i < 2 * args.count else [] for i, x in enumerate(index)
]
train_index = [x[:-1] if i < args.count else x for i, x in enumerate(index)]
dataset = env.MultispeakerDataset(train_index, data_path)
if model_type == 'vqvae':
model_fn = lambda dataset: vqvae.Model(rnn_dims=896,
fc_dims=896,
global_decoder_cond_dims=dataset[0][
0].shape[0],
upsample_factors=(4, 4, 4),
normalize_vq=True,
noise_x=True,
noise_y=True).cuda()
dataset_type = 'multi'
elif model_type == 'wavernn':
model_fn = lambda dataset: wr.Model(rnn_dims=896,
fc_dims=896,
pad=2,
upsample_factors=(4, 4, 4),
feat_dims=80).cuda()
dataset_type = 'single'
elif model_type == 'nc':
model_fn = lambda dataset: nc.Model(rnn_dims=896, fc_dims=896).cuda()
dataset_type = 'single'
else:
if args.float:
use_half = False
elif args.half:
use_half = True
else:
use_half = False
model_type = args.model or 'vqvae'
model_name = f'{model_type}.43.upconv'
if model_type == 'vqvae':
model_fn = lambda dataset: vqvae.Model(
rnn_dims=config.rnn_dims,
fc_dims=config.fc_dims,
global_decoder_cond_dims=dataset.num_speakers(),
upsample_factors=config.upsample_factors,
normalize_vq=True,
noise_x=True,
noise_y=True).cuda()
dataset_type = 'multi'
elif model_type == 'wavernn':
model_fn = lambda dataset: wr.Model(rnn_dims=config.rnn_dims,
fc_dims=config.fc_dims,
pad=2,
upsample_factors=config.
upsample_factors,
feat_dims=80).cuda()
dataset_type = 'single'
elif model_type == 'nc':
model_fn = lambda dataset: nc.Model(rnn_dims=896, fc_dims=896).cuda()
dataset_type = 'single'
else:
use_half = False
DEVICE = "cpu"
model_type = args.model or 'vqvae'
#model_name = f'{model_type}.43.upconv'
model_name = f'{model_type}.43.upconv'
if model_type == 'vqvae':
model_fn = lambda dataset: vqvae.Model(rnn_dims=896,
fc_dims=896,
global_decoder_cond_dims=dataset.
num_speakers(),
upsample_factors=(4, 4, 4),
normalize_vq=True,
noise_x=True,
noise_y=True,
DEVICE=DEVICE).to(DEVICE)
dataset_type = 'multi'
elif model_type == 'wavernn':
raise ValueError("NYI wavernn")
model_fn = lambda dataset: wr.Model(rnn_dims=896,
fc_dims=896,
pad=2,
upsample_factors=(4, 4, 4),
feat_dims=80,
DEVICE=DEVICE).to(DEVICE)
dataset_type = 'single'
elif model_type == 'nc':