def _sample(zs, labels, sampling_kwargs, priors, sample_levels, hps):
alignments = None
for level in reversed(sample_levels):
prior = priors[level]
prior.cuda()
empty_cache()
# Set correct total_length, hop_length, labels and sampling_kwargs for level
assert hps.sample_length % prior.raw_to_tokens == 0, f"Expected sample_length {hps.sample_length} to be multiple of {prior.raw_to_tokens}"
total_length = hps.sample_length//prior.raw_to_tokens
hop_length = int(hps.hop_fraction[level]*prior.n_ctx)
zs = sample_level(zs, labels[level], sampling_kwargs[level], level, prior, total_length, hop_length, hps)
prior.cpu()
empty_cache()
# Decode sample
x = prior.decode(zs[level:], start_level=level, bs_chunks=zs[level].shape[0])
logdir = f"{hps.name}/level_{level}"
if not os.path.exists(logdir):
os.makedirs(logdir)
t.save(dict(zs=zs, labels=labels, sampling_kwargs=sampling_kwargs, x=x), f"{logdir}/data.pth.tar")
save_wav(logdir, x, hps.sr)
if alignments is None and priors[-1] is not None and priors[-1].n_tokens > 0:
alignments = get_alignment(x, zs, labels[-1], priors[-1], sampling_kwargs[-1]['fp16'], hps)
save_html(logdir, x, zs, labels[-1], alignments, hps)
return zs
def _sample(zs, labels_1, labels_2, sampling_kwargs, priors, sample_levels,
hps):
alignments = None
for level in reversed(sample_levels):
prior = priors[level]
prior.cuda()
empty_cache()
# Set correct total_length, hop_length, labels and sampling_kwargs for level
assert hps.sample_length % prior.raw_to_tokens == 0, f"Expected sample_length {hps.sample_length} to be multiple of {prior.raw_to_tokens}"
total_length = hps.sample_length // prior.raw_to_tokens
hop_length = int(hps.hop_fraction[level] * prior.n_ctx)
zs = sample_level(zs, labels_1[level], labels_2[level],
sampling_kwargs[level], level, prior, total_length,
hop_length, hps)
prior.cpu()
empty_cache()
# Decode sample
x = prior.decode(zs[level:],
start_level=level,
bs_chunks=zs[level].shape[0])
if dist.get_world_size() > 1:
logdir = f"{hps.name}_rank_{dist.get_rank()}/level_{level}"
else:
logdir = f"{hps.name}/level_{level}"
if not os.path.exists(logdir):
os.makedirs(logdir)
t.save(
dict(zs=zs, labels=labels_1, sampling_kwargs=sampling_kwargs, x=x),
f"{logdir}/data.pth.tar")
save_wav(logdir, x, hps.sr)
if alignments is None and priors[
-1] is not None and priors[-1].n_tokens > 0 and not isinstance(
priors[-1].labeller, EmptyLabeller):
try:
labels_1[-1], priors[-1], sampling_kwargs[-1]['fp16']
except:
import ipdb
ipdb.set_trace()
alignments = get_alignment(x, zs, labels_1[-1], priors[-1],
sampling_kwargs[-1]['fp16'], hps)
# don't care
# save_html(logdir, x, zs, labels_1[-1], alignments, hps)
return zs