def sample_single_window(zs, labels, sampling_kwargs, level, prior, start,
hps):
n_samples = hps.n_samples
n_ctx = prior.n_ctx
end = start + n_ctx
# get z already sampled at current level
z = zs[level][:, start:end]
if 'sample_tokens' in sampling_kwargs:
# Support sampling a window shorter than n_ctx
sample_tokens = sampling_kwargs['sample_tokens']
else:
sample_tokens = (end - start)
conditioning_tokens, new_tokens = z.shape[1], sample_tokens - z.shape[1]
print_once(
f"Sampling {sample_tokens} tokens for [{start},{start+sample_tokens}]. Conditioning on {conditioning_tokens} tokens"
)
if new_tokens <= 0:
# Nothing new to sample
return zs
# get z_conds from level above
z_conds = prior.get_z_conds(zs, start, end)
# set y offset, sample_length and lyrics tokens
y = prior.get_y(labels, start)
empty_cache()
max_batch_size = sampling_kwargs['max_batch_size']
del sampling_kwargs['max_batch_size']
z_list = split_batch(z, n_samples, max_batch_size)
z_conds_list = split_batch(z_conds, n_samples, max_batch_size)
y_list = split_batch(y, n_samples, max_batch_size)
z_samples = []
for z_i, z_conds_i, y_i in zip(z_list, z_conds_list, y_list):
z_samples_i = prior.sample(n_samples=z_i.shape[0],
z=z_i,
z_conds=z_conds_i,
y=y_i,
**sampling_kwargs)
z_samples.append(z_samples_i)
z = t.cat(z_samples, dim=0)
sampling_kwargs['max_batch_size'] = max_batch_size
# Update z with new sample
z_new = z[:, -new_tokens:]
zs[level] = t.cat([zs[level], z_new], dim=1)
return zs
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])
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, 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):
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 get_alignment(x, zs, labels, prior, fp16, hps):
level = hps.levels - 1 # Top level used
n_ctx, n_tokens = prior.n_ctx, prior.n_tokens
z = zs[level]
bs, total_length = z.shape[0], z.shape[1]
if total_length < n_ctx:
padding_length = n_ctx - total_length
z = t.cat([z, t.zeros(bs, n_ctx - total_length, dtype=z.dtype, device=z.device)], dim=1)
total_length = z.shape[1]
else:
padding_length = 0
hop_length = int(hps.hop_fraction[level]*prior.n_ctx)
n_head = prior.prior.transformer.n_head
alignment_head, alignment_layer = prior.alignment_head, prior.alignment_layer
attn_layers = set([alignment_layer])
alignment_hops = {}
indices_hops = {}
prior.cuda()
empty_cache()
for start in get_starts(total_length, n_ctx, hop_length):
end = start + n_ctx
# set y offset, sample_length and lyrics tokens
y, indices_hop = prior.get_y(labels, start, get_indices=True)
assert len(indices_hop) == bs
for indices in indices_hop:
assert len(indices) == n_tokens
z_bs = t.chunk(z, bs, dim=0)
y_bs = t.chunk(y, bs, dim=0)
w_hops = []
for z_i, y_i in zip(z_bs, y_bs):
w_hop = prior.z_forward(z_i[:,start:end], [], y_i, fp16=fp16, get_attn_weights=attn_layers)
assert len(w_hop) == 1
w_hops.append(w_hop[0][:, alignment_head])
del w_hop
w = t.cat(w_hops, dim=0)
del w_hops
assert_shape(w, (bs, n_ctx, n_tokens))
alignment_hop = w.float().cpu().numpy()
assert_shape(alignment_hop, (bs, n_ctx, n_tokens))
del w
# alignment_hop has shape (bs, n_ctx, n_tokens)
# indices_hop is a list of len=bs, each entry of len hps.n_tokens
indices_hops[start] = indices_hop
alignment_hops[start] = alignment_hop
prior.cpu()
empty_cache()
# Combine attn for each hop into attn for full range
# Use indices to place them into correct place for corresponding source tokens
alignments = []
for item in range(bs):
# Note each item has different length lyrics
full_tokens = labels['info'][item]['full_tokens']
alignment = np.zeros((total_length, len(full_tokens) + 1))
for start in reversed(get_starts(total_length, n_ctx, hop_length)):
end = start + n_ctx
alignment_hop = alignment_hops[start][item]
indices = indices_hops[start][item]
assert len(indices) == n_tokens
assert alignment_hop.shape == (n_ctx, n_tokens)
alignment[start:end,indices] = alignment_hop
alignment = alignment[:total_length - padding_length,:-1] # remove token padding, and last lyric index
alignments.append(alignment)
return alignments
def primed_sample(self,
n_samples,
x,
x_cond=None,
y_cond=None,
encoder_kv=None,
fp16=False,
temp=1.0,
top_k=0,
top_p=0.0,
get_preds=False,
chunk_size=None,
sample_tokens=None):
assert self.training == False
if sample_tokens is None: sample_tokens = self.input_dims
# Preprocess.
with t.no_grad():
x = self.preprocess(x)
assert isinstance(x, t.cuda.LongTensor)
assert (0 <= x).all() and (x < self.bins).all()
assert x.shape[0] == n_samples
xs = t.split(x, 1, dim=1)
xs = list(xs)
assert len(xs) < sample_tokens
N, D = n_samples, self.input_dims
if self.y_cond:
assert y_cond is not None
assert y_cond.shape == (N, 1, self.width)
else:
assert y_cond is None
if self.x_cond:
assert x_cond is not None
assert x_cond.shape == (N, D, self.width) or x_cond.shape == (
N, 1, self.width
), f"Got {x_cond.shape}, expected ({N}, {D}/{1}, {self.width})"
else:
assert x_cond is None
x_cond = t.zeros((N, 1, self.width), dtype=t.float).cuda()
with t.no_grad():
if get_preds:
preds = []
# Fill up key/value cache for past context by runing forward pass.
# We do so in chunks instead of doing the whole past in one forward pass to reduce max memory usage.
if chunk_size is None:
chunk_size = len(xs)
#assert len(xs) % chunk_size == 0, f'expected {len(xs)} to be divisible by {chunk_size}'
chunk_sizes = split_chunks(len(xs), chunk_size)
x_primes = []
start = 0
x = None
for current_chunk_size in get_range(chunk_sizes):
xs_prime, conds_prime = [], []
for sample_t in range(start, start + current_chunk_size):
x_prime, cond_prime = self.get_emb(sample_t, n_samples, x,
x_cond, y_cond)
x = xs[sample_t]
xs_prime.append(x_prime)
conds_prime.append(cond_prime)
start = start + current_chunk_size
x_prime, cond_prime = t.cat(xs_prime,
dim=1), t.cat(conds_prime, dim=1)
assert x_prime.shape == (n_samples, current_chunk_size,
self.width)
assert cond_prime.shape == (n_samples, current_chunk_size,
self.width)
del xs_prime
del conds_prime
if not get_preds:
del cond_prime
x_prime = self.transformer(x_prime,
encoder_kv=encoder_kv,
sample=True,
fp16=fp16)
if get_preds:
if self.add_cond_after_transformer:
x_prime = x_prime + cond_prime
assert x_prime.shape == (n_samples, current_chunk_size,
self.width)
del cond_prime
x_primes.append(x_prime)
else:
del x_prime
if get_preds:
x_prime = t.cat(x_primes, dim=1)
assert x_prime.shape == (n_samples, len(xs), self.width)
x_prime = self.x_out(x_prime) # Predictions
preds.append(x_prime)
empty_cache()
self.transformer.check_cache(n_samples, len(xs), fp16)
x = xs[-1]
assert x.shape == (n_samples, 1)
empty_cache()
for sample_t in get_range(range(len(xs), sample_tokens)):
x, cond = self.get_emb(sample_t, n_samples, x, x_cond, y_cond)
self.transformer.check_cache(n_samples, sample_t, fp16)
x = self.transformer(x,
encoder_kv=encoder_kv,
sample=True,
fp16=fp16) # Transformer
if self.add_cond_after_transformer:
x = x + cond
assert x.shape == (n_samples, 1, self.width)
x = self.x_out(x) # Predictions
if get_preds:
preds.append(x)
# Adjust logits
x = x / temp
x = filter_logits(x, top_k=top_k, top_p=top_p)
x = t.distributions.Categorical(
logits=x).sample() # Sample and replace x
assert x.shape == (n_samples, 1)
xs.append(x.clone())
del x
self.transformer.del_cache()
x = t.cat(xs, dim=1)
if get_preds:
preds = t.cat(preds, dim=1)
x = self.postprocess(x, sample_tokens)
if get_preds:
return x, preds
else:
return x