Spaces:
Build error
Build error
import io | |
import os | |
import random | |
import time | |
from glob import glob | |
import IPython | |
import librosa | |
import matplotlib.pyplot as plt | |
import numpy as np | |
import soundfile as sf | |
import tensorflow as tf | |
import tensorflow_io as tfio | |
from tensorflow.keras import mixed_precision | |
from tensorflow.keras.optimizers import Adam | |
from tensorflow.python.framework import random_seed | |
from tqdm import tqdm | |
import gradio as gr | |
from scipy.io.wavfile import write as write_wav | |
class Utils_functions: | |
def __init__(self, args): | |
self.args = args | |
melmat = tf.signal.linear_to_mel_weight_matrix( | |
num_mel_bins=args.mel_bins, | |
num_spectrogram_bins=(4 * args.hop) // 2 + 1, | |
sample_rate=args.sr, | |
lower_edge_hertz=0.0, | |
upper_edge_hertz=args.sr // 2, | |
) | |
mel_f = tf.convert_to_tensor(librosa.mel_frequencies(n_mels=args.mel_bins + 2, fmin=0.0, fmax=args.sr // 2)) | |
enorm = tf.cast( | |
tf.expand_dims(tf.constant(2.0 / (mel_f[2 : args.mel_bins + 2] - mel_f[: args.mel_bins])), 0,), tf.float32, | |
) | |
melmat = tf.multiply(melmat, enorm) | |
melmat = tf.divide(melmat, tf.reduce_sum(melmat, axis=0)) | |
self.melmat = tf.where(tf.math.is_nan(melmat), tf.zeros_like(melmat), melmat) | |
with np.errstate(divide="ignore", invalid="ignore"): | |
self.melmatinv = tf.constant(np.nan_to_num(np.divide(melmat.numpy().T, np.sum(melmat.numpy(), axis=1))).T) | |
def conc_tog_specphase(self, S, P): | |
S = tf.cast(S, tf.float32) | |
P = tf.cast(P, tf.float32) | |
S = self.denormalize(S, clip=False) | |
S = tf.math.sqrt(self.db2power(S) + 1e-7) | |
P = P * np.pi | |
Sls = tf.split(S, S.shape[0], 0) | |
S = tf.squeeze(tf.concat(Sls, 1), 0) | |
Pls = tf.split(P, P.shape[0], 0) | |
P = tf.squeeze(tf.concat(Pls, 1), 0) | |
SP = tf.cast(S, tf.complex64) * tf.math.exp(1j * tf.cast(P, tf.complex64)) | |
wv = tf.signal.inverse_stft( | |
SP, | |
4 * self.args.hop, | |
self.args.hop, | |
fft_length=4 * self.args.hop, | |
window_fn=tf.signal.inverse_stft_window_fn(self.args.hop), | |
) | |
return np.squeeze(wv) | |
def _tf_log10(self, x): | |
numerator = tf.math.log(x) | |
denominator = tf.math.log(tf.constant(10, dtype=numerator.dtype)) | |
return numerator / denominator | |
def normalize(self, S, clip=False): | |
S = (S - self.args.mu_rescale) / self.args.sigma_rescale | |
if clip: | |
S = tf.clip_by_value(S, -1.0, 1.0) | |
return S | |
def normalize_rel(self, S): | |
S = S - tf.math.reduce_min(S + 1e-7) | |
S = (S / (tf.math.reduce_max(S + 1e-7) + 1e-7)) + 1e-7 | |
return S | |
def denormalize(self, S, clip=False): | |
if clip: | |
S = tf.clip_by_value(S, -1.0, 1.0) | |
return (S * self.args.sigma_rescale) + self.args.mu_rescale | |
def amp2db(self, x): | |
return 20 * self._tf_log10(tf.clip_by_value(tf.abs(x), 1e-5, 1e100)) | |
def db2amp(self, x): | |
return tf.pow(tf.ones(tf.shape(x)) * 10.0, x * 0.05) | |
def power2db(self, power, ref_value=1.0, amin=1e-10, top_db=None, norm=False): | |
log_spec = 10.0 * self._tf_log10(tf.maximum(amin, power)) | |
log_spec -= 10.0 * self._tf_log10(tf.maximum(amin, ref_value)) | |
if top_db is not None: | |
log_spec = tf.maximum(log_spec, tf.reduce_max(log_spec) - top_db) | |
return log_spec | |
def power2db_batch(self, power, ref_value=1.0, amin=1e-10, top_db=None, norm=False): | |
log_spec = 10.0 * self._tf_log10(tf.maximum(amin, power)) | |
log_spec -= 10.0 * self._tf_log10(tf.maximum(amin, ref_value)) | |
if top_db is not None: | |
log_spec = tf.maximum(log_spec, tf.reduce_max(log_spec, [-2, -1], keepdims=True) - top_db) | |
return log_spec | |
def db2power(self, S_db, ref=1.0): | |
return ref * tf.math.pow(10.0, 0.1 * S_db) | |
def wv2mel(self, wv, topdb=80.0): | |
X = tf.signal.stft( | |
wv, | |
frame_length=4 * self.args.hop, | |
frame_step=self.args.hop, | |
fft_length=4 * self.args.hop, | |
window_fn=tf.signal.hann_window, | |
pad_end=False, | |
) | |
S = self.normalize(self.power2db(tf.abs(X) ** 2, top_db=topdb) - self.args.ref_level_db) | |
SM = tf.tensordot(S, self.melmat, 1) | |
return SM | |
def mel2spec(self, SM): | |
return tf.tensordot(SM, tf.transpose(self.melmatinv), 1) | |
def spec2mel(self, S): | |
return tf.tensordot(S, self.melmat, 1) | |
def wv2spec(self, wv, hop_size=256, fac=4): | |
X = tf.signal.stft( | |
wv, | |
frame_length=fac * hop_size, | |
frame_step=hop_size, | |
fft_length=fac * hop_size, | |
window_fn=tf.signal.hann_window, | |
pad_end=False, | |
) | |
return self.normalize(self.power2db(tf.abs(X) ** 2, top_db=None)) | |
def wv2spec_hop(self, wv, topdb=80.0, hopsize=256): | |
X = tf.signal.stft( | |
wv, | |
frame_length=4 * hopsize, | |
frame_step=hopsize, | |
fft_length=4 * hopsize, | |
window_fn=tf.signal.hann_window, | |
pad_end=False, | |
) | |
S = self.normalize(self.power2db(tf.abs(X) ** 2, top_db=topdb)) | |
return tf.tensordot(S, self.melmat, 1) | |
def distribute(self, x, model, bs=64, dual_out=False): | |
outls = [] | |
if isinstance(x, list): | |
bdim = x[0].shape[0] | |
for i in range(((bdim - 2) // bs) + 1): | |
outls.append(model([el[i * bs : i * bs + bs] for el in x], training=False)) | |
else: | |
bdim = x.shape[0] | |
for i in range(((bdim - 2) // bs) + 1): | |
outls.append(model(x[i * bs : i * bs + bs], training=False)) | |
if dual_out: | |
return ( | |
np.concatenate([outls[k][0] for k in range(len(outls))], 0), | |
np.concatenate([outls[k][1] for k in range(len(outls))], 0), | |
) | |
else: | |
return np.concatenate(outls, 0) | |
def distribute_enc(self, x, model, bs=64): | |
outls = [] | |
if isinstance(x, list): | |
bdim = x[0].shape[0] | |
for i in range(((bdim - 2) // bs) + 1): | |
res = model([el[i * bs : i * bs + bs] for el in x], training=False) | |
resls = tf.split(res, self.args.shape // self.args.window, 0) | |
res = tf.concat(resls, -2) | |
outls.append(res) | |
else: | |
bdim = x.shape[0] | |
for i in range(((bdim - 2) // bs) + 1): | |
res = model(x[i * bs : i * bs + bs], training=False) | |
resls = tf.split(res, self.args.shape // self.args.window, 0) | |
res = tf.concat(resls, -2) | |
outls.append(res) | |
return np.concatenate(outls, 0) | |
def distribute_dec(self, x, model, bs=64): | |
outls = [] | |
bdim = x.shape[0] | |
for i in range(((bdim - 2) // bs) + 1): | |
inp = x[i * bs : i * bs + bs] | |
inpls = tf.split(inp, 2, -2) | |
inp = tf.concat(inpls, 0) | |
res = model(inp, training=False) | |
outls.append(res) | |
return ( | |
np.concatenate([outls[k][0] for k in range(len(outls))], 0), | |
np.concatenate([outls[k][1] for k in range(len(outls))], 0), | |
) | |
def distribute_dec2(self, x, model, bs=64): | |
outls = [] | |
bdim = x.shape[0] | |
for i in range(((bdim - 2) // bs) + 1): | |
inp1 = x[i * bs : i * bs + bs] | |
inpls = tf.split(inp1, 2, -2) | |
inp1 = tf.concat(inpls, 0) | |
outls.append(model(inp1, training=False)) | |
return np.concatenate(outls, 0) | |
def get_noise_interp(self): | |
noiseg = tf.random.normal([1, 64], dtype=tf.float32) | |
noisel = tf.concat([tf.random.normal([1, 64], dtype=tf.float32), noiseg], -1) | |
noisec = tf.concat([tf.random.normal([1, 64], dtype=tf.float32), noiseg], -1) | |
noiser = tf.concat([tf.random.normal([1, 64], dtype=tf.float32), noiseg], -1) | |
rl = tf.linspace(noisel, noisec, self.args.latlen + 1, axis=-2)[:, :-1, :] | |
rr = tf.linspace(noisec, noiser, self.args.latlen + 1, axis=-2) | |
noisetot = tf.concat([rl, rr], -2) | |
return tf.image.random_crop(noisetot, [1, self.args.latlen, 64 + 64]) | |
def generate_example_stereo(self, models_ls): | |
(critic, gen, enc, dec, enc2, dec2, critic_rec, gen_ema, [opt_dec, opt_disc],) = models_ls | |
abb = gen_ema(self.get_noise_interp(), training=False) | |
abbls = tf.split(abb, abb.shape[-2] // 16, -2) | |
abb = tf.concat(abbls, 0) | |
chls = [] | |
for channel in range(2): | |
ab = self.distribute_dec2( | |
abb[:, :, :, channel * self.args.latdepth : channel * self.args.latdepth + self.args.latdepth,], dec2, | |
) | |
abls = tf.split(ab, ab.shape[-2] // self.args.shape, -2) | |
ab = tf.concat(abls, 0) | |
ab_m, ab_p = self.distribute_dec(ab, dec) | |
wv = self.conc_tog_specphase(ab_m, ab_p) | |
chls.append(wv) | |
return np.stack(chls, -1) | |
# Save in training loop | |
def save_test_image_full(self, path, models_ls=None): | |
abwv = self.generate_example_stereo(models_ls) | |
abwv2 = self.generate_example_stereo(models_ls) | |
abwv3 = self.generate_example_stereo(models_ls) | |
abwv4 = self.generate_example_stereo(models_ls) | |
# IPython.display.display( | |
# IPython.display.Audio(np.squeeze(np.transpose(abwv)), rate=self.args.sr) | |
# ) | |
# IPython.display.display( | |
# IPython.display.Audio(np.squeeze(np.transpose(abwv2)), rate=self.args.sr) | |
# ) | |
# IPython.display.display( | |
# IPython.display.Audio(np.squeeze(np.transpose(abwv3)), rate=self.args.sr) | |
# ) | |
# IPython.display.display( | |
# IPython.display.Audio(np.squeeze(np.transpose(abwv4)), rate=self.args.sr) | |
# ) | |
write_wav(f"{path}/out1.wav", self.args.sr, np.squeeze(abwv)) | |
write_wav(f"{path}/out2.wav", self.args.sr, np.squeeze(abwv2)) | |
write_wav(f"{path}/out3.wav", self.args.sr, np.squeeze(abwv3)) | |
write_wav(f"{path}/out4.wav", self.args.sr, np.squeeze(abwv4)) | |
fig, axs = plt.subplots(nrows=4, ncols=1, figsize=(20, 20)) | |
axs[0].imshow( | |
np.flip(np.array(tf.transpose(self.wv2spec_hop((abwv[:, 0] + abwv[:, 1]) / 2.0, 80.0, 256), [1, 0],)), -2,), | |
cmap=None, | |
) | |
axs[0].axis("off") | |
axs[0].set_title("Generated1") | |
axs[1].imshow( | |
np.flip( | |
np.array(tf.transpose(self.wv2spec_hop((abwv2[:, 0] + abwv2[:, 1]) / 2.0, 80.0, 256), [1, 0],)), -2, | |
), | |
cmap=None, | |
) | |
axs[1].axis("off") | |
axs[1].set_title("Generated2") | |
axs[2].imshow( | |
np.flip( | |
np.array(tf.transpose(self.wv2spec_hop((abwv3[:, 0] + abwv3[:, 1]) / 2.0, 80.0, 256), [1, 0],)), -2, | |
), | |
cmap=None, | |
) | |
axs[2].axis("off") | |
axs[2].set_title("Generated3") | |
axs[3].imshow( | |
np.flip( | |
np.array(tf.transpose(self.wv2spec_hop((abwv4[:, 0] + abwv4[:, 1]) / 2.0, 80.0, 256), [1, 0],)), -2, | |
), | |
cmap=None, | |
) | |
axs[3].axis("off") | |
axs[3].set_title("Generated4") | |
# plt.show() | |
plt.savefig(f"{path}/output.png") | |
# Save in training loop | |
def save_end( | |
self, epoch, gloss, closs, mloss, models_ls=None, n_save=3, save_path="checkpoints", | |
): | |
(critic, gen, enc, dec, enc2, dec2, critic_rec, gen_ema, [opt_dec, opt_disc],) = models_ls | |
if epoch % n_save == 0: | |
print("Saving...") | |
path = f"{save_path}/MUSIKA!_-{str(gloss)[:9]}-{str(closs)[:9]}-{str(mloss)[:9]}" | |
os.mkdir(path) | |
critic.save_weights(path + "/critic.h5") | |
critic_rec.save_weights(path + "/critic_rec.h5") | |
gen.save_weights(path + "/gen.h5") | |
gen_ema.save_weights(path + "/gen_ema.h5") | |
# enc.save_weights(path + "/enc.h5") | |
# dec.save_weights(path + "/dec.h5") | |
# enc2.save_weights(path + "/enc2.h5") | |
# dec2.save_weights(path + "/dec2.h5") | |
np.save(path + "/opt_dec.npy", opt_dec.get_weights()) | |
np.save(path + "/opt_disc.npy", opt_disc.get_weights()) | |
self.save_test_image_full(path, models_ls=models_ls) | |
def truncated_normal(self, shape, bound=2.0, dtype=tf.float32): | |
seed1, seed2 = random_seed.get_seed(tf.random.uniform((), tf.int32.min, tf.int32.max, dtype=tf.int32)) | |
return tf.random.stateless_parameterized_truncated_normal(shape, [seed1, seed2], 0.0, 1.0, -bound, bound) | |
def distribute_gen(self, x, model, bs=64): | |
outls = [] | |
bdim = x.shape[0] | |
if bdim == 1: | |
bdim = 2 | |
for i in range(((bdim - 2) // bs) + 1): | |
outls.append(model(x[i * bs : i * bs + bs], training=False)) | |
return np.concatenate(outls, 0) | |
def get_noise_interp_multi(self, fac=1, var=2.0): | |
noiseg = self.truncated_normal([1, 64], var, dtype=tf.float32) | |
if var < 1.75: | |
var = 1.75 | |
noisels = [ | |
tf.concat([self.truncated_normal([1, 64], var, dtype=tf.float32), noiseg], -1) for i in range(2 + (fac - 1)) | |
] | |
rls = [ | |
tf.linspace(noisels[k], noisels[k + 1], self.args.latlen + 1, axis=-2)[:, :-1, :] | |
for k in range(len(noisels) - 1) | |
] | |
return tf.concat(rls, 0) | |
def stfunc(self, genre, z, var, models_ls_techno, models_ls_classical): | |
( | |
critic, | |
gen, | |
enc, | |
dec_techno, | |
enc2, | |
dec2_techno, | |
critic_rec, | |
gen_ema_techno, | |
[opt_dec, opt_disc], | |
) = models_ls_techno | |
( | |
critic, | |
gen, | |
enc, | |
dec_classical, | |
enc2, | |
dec2_classical, | |
critic_rec, | |
gen_ema_classical, | |
[opt_dec, opt_disc], | |
) = models_ls_classical | |
var = 0.01 + (3.5 * (var / 100.0)) | |
if z == 0: | |
fac = 1 | |
elif z == 1: | |
fac = 5 | |
else: | |
fac = 10 | |
if genre == 0: | |
dec = dec_techno | |
dec2 = dec2_techno | |
gen_ema = gen_ema_techno | |
else: | |
dec = dec_classical | |
dec2 = dec2_classical | |
gen_ema = gen_ema_classical | |
bef = time.time() | |
ab = self.distribute_gen(self.get_noise_interp_multi(fac, var), gen_ema) | |
abls = tf.split(ab, ab.shape[0], 0) | |
ab = tf.concat(abls, -2) | |
abls = tf.split(ab, ab.shape[-2] // 16, -2) | |
abi = tf.concat(abls, 0) | |
chls = [] | |
for channel in range(2): | |
ab = self.distribute_dec2( | |
abi[:, :, :, channel * self.args.latdepth : channel * self.args.latdepth + self.args.latdepth], | |
dec2, | |
bs=128, | |
) | |
# abls = tf.split(ab, ab.shape[-2] // (self.args.shape // 2), -2) | |
abls = tf.split(ab, ab.shape[-2] // self.args.shape, -2) | |
ab = tf.concat(abls, 0) | |
ab_m, ab_p = self.distribute_dec(ab, dec, bs=128) | |
abwv = self.conc_tog_specphase(ab_m, ab_p) | |
chls.append(abwv) | |
print( | |
f"Time for complete generation pipeline: {time.time()-bef} s {int(np.round((fac*23.)/(time.time()-bef)))}x faster than Real Time!" | |
) | |
abwvc = np.clip(np.squeeze(np.stack(chls, -1)), -1.0, 1.0) | |
spec = np.flip( | |
np.array( | |
tf.transpose( | |
self.wv2spec_hop((abwvc[: 23 * self.args.sr, 0] + abwvc[: 23 * self.args.sr, 1]) / 2.0, 80.0, 256), | |
[1, 0], | |
) | |
), | |
-2, | |
) | |
return ( | |
spec, | |
(self.args.sr, np.int16(abwvc * 32767.0)), | |
) | |
def render_gradio(self, models_ls_techno, models_ls_classical, train=True): | |
article_text = "Original work by Marco Pasini ([Twitter](https://twitter.com/marco_ppasini)) and Jan Schlüter at Johannes Kepler Universität Linz." | |
def gradio_func(x, y, z): | |
return self.stfunc(x, y, z, models_ls_techno, models_ls_classical) | |
iface = gr.Interface( | |
fn=gradio_func, | |
inputs=[ | |
gr.inputs.Radio( | |
choices=["Techno/Experimental", "Classical"], | |
type="index", | |
default="Classical", | |
label="Music Genre to Generate", | |
), | |
gr.inputs.Radio( | |
choices=["23s", "1m 58s", "3m 57s"], type="index", default="1m 58s", label="Generated Music Length", | |
), | |
gr.inputs.Slider( | |
minimum=0, | |
maximum=100, | |
step=1, | |
default=25, | |
label="Stability[left]/Variety[right] Tradeoff (Truncation Trick)", | |
), | |
], | |
outputs=[ | |
gr.outputs.Image(label="Log-MelSpectrogram of Generated Audio (first 23 sec)"), | |
gr.outputs.Audio(type="numpy", label="Generated Audio"), | |
], | |
allow_screenshot=False, | |
title="musika!", | |
description="Blazingly Fast Stereo Waveform Music Generation of Arbitrary Length. Be patient and enjoy the weirdness!", | |
article=article_text, | |
layout="vertical", | |
theme="huggingface", | |
) | |
print("--------------------------------") | |
print("--------------------------------") | |
print("--------------------------------") | |
print("--------------------------------") | |
print("--------------------------------") | |
print("CLICK ON LINK BELOW TO OPEN GRADIO INTERFACE") | |
if train: | |
iface.launch(prevent_thread_lock=True) | |
else: | |
iface.launch(enable_queue=True) | |
# iface.launch(share=True, enable_queue=True) | |
print("--------------------------------") | |
print("--------------------------------") | |
print("--------------------------------") | |
print("--------------------------------") | |
print("--------------------------------") | |