import os import time import datetime from glob import glob from tqdm import tqdm import librosa import matplotlib.pyplot as plt import numpy as np import tensorflow as tf from tensorflow.python.framework import random_seed 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) // 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 rand_channel_swap(self, x): s_l, s_r = tf.split(x, 2, -1) if tf.random.uniform((), dtype=tf.float32) > 0.5: sl = s_l sr = s_r else: sl = s_r sr = s_l return tf.concat([sl, sr], -1) def distribute(self, x, model, bs=32, 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=32): 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 tf.concat(outls, 0) def distribute_dec(self, x, model, bs=32): 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=32): 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 tf.concat(outls, 0) def center_coordinate( self, x ): # allows to have sequences with even number length with anchor in the middle of the sequence return tf.reduce_mean(tf.stack([x, tf.roll(x, -1, -2)], 0), 0)[:, :-1, :] # hardcoded! need to figure out how to generalize it without breaking jit compiling def crop_coordinate( self, x ): # randomly crops a conditioning sequence such that the anchor vector is at center of generator receptive field (maximum context is provided to the generator) fac = tf.random.uniform((), 0, self.args.coordlen // (self.args.latlen // 2), dtype=tf.int32) if fac == 0: return tf.reshape( x[ :, (self.args.latlen // 4) + 0 * (self.args.latlen // 2) : (self.args.latlen // 4) + 0 * (self.args.latlen // 2) + self.args.latlen, :, ], [-1, self.args.latlen, x.shape[-1]], ) elif fac == 1: return tf.reshape( x[ :, (self.args.latlen // 4) + 1 * (self.args.latlen // 2) : (self.args.latlen // 4) + 1 * (self.args.latlen // 2) + self.args.latlen, :, ], [-1, self.args.latlen, x.shape[-1]], ) else: return tf.reshape( x[ :, (self.args.latlen // 4) + 2 * (self.args.latlen // 2) : (self.args.latlen // 4) + 2 * (self.args.latlen // 2) + self.args.latlen, :, ], [-1, self.args.latlen, x.shape[-1]], ) def update_switch(self, switch, ca, cab, learning_rate_switch=0.0001, stable_point=0.9): cert = tf.math.minimum(tf.math.maximum(tf.reduce_mean(ca) - tf.reduce_mean(cab), 0.0), 2.0) / 2.0 if cert > stable_point: switch_new = switch - learning_rate_switch else: switch_new = switch + learning_rate_switch return tf.math.maximum(tf.math.minimum(switch_new, 0.0), -1.0) def get_noise_interp(self): noiseg = tf.random.normal([1, 64], dtype=tf.float32) noisel = tf.concat([tf.random.normal([1, self.args.coorddepth], dtype=tf.float32), noiseg], -1) noisec = tf.concat([tf.random.normal([1, self.args.coorddepth], dtype=tf.float32), noiseg], -1) noiser = tf.concat([tf.random.normal([1, self.args.coorddepth], dtype=tf.float32), noiseg], -1) rl = tf.linspace(noisel, noisec, self.args.coordlen + 1, axis=-2)[:, :-1, :] rr = tf.linspace(noisec, noiser, self.args.coordlen + 1, axis=-2) noisetot = tf.concat([rl, rr], -2) noisetot = self.center_coordinate(noisetot) return self.crop_coordinate(noisetot) def generate_example_stereo(self, models_ls): ( critic, gen, enc, dec, enc2, dec2, gen_ema, [opt_dec, opt_disc], switch, ) = models_ls abb = gen_ema(self.get_noise_interp(), training=False) abbls = tf.split(abb, abb.shape[-2] // 8, -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, self.args.hop * 2), [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, self.args.hop * 2), [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, self.args.hop * 2), [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, self.args.hop * 2), [1, 0], ) ), -2, ), cmap=None, ) axs[3].axis("off") axs[3].set_title("Generated4") # plt.show() plt.savefig(f"{path}/output.png") plt.close() def save_end( self, epoch, gloss, closs, mloss, models_ls=None, n_save=3, save_path="checkpoints", ): (critic, gen, enc, dec, enc2, dec2, gen_ema, [opt_dec, opt_disc], switch) = models_ls if epoch % n_save == 0: print("Saving...") path = f"{save_path}/MUSIKA_iterations-{((epoch+1)*self.args.totsamples)//(self.args.bs*1000)}k_losses-{str(gloss)[:9]}-{str(closs)[:9]}-{str(mloss)[:9]}" os.mkdir(path) critic.save_weights(path + "/critic.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()) np.save(path + "/switch.npy", switch.numpy()) 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=32): 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 tf.concat(outls, 0) def generate_waveform(self, inp, gen_ema, dec, dec2, batch_size=64): ab = self.distribute_gen(inp, gen_ema, bs=batch_size) abls = tf.split(ab, ab.shape[0], 0) ab = tf.concat(abls, -2) abls = tf.split(ab, ab.shape[-2] // 8, -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=batch_size, ) 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=batch_size) abwv = self.conc_tog_specphase(ab_m, ab_p) chls.append(abwv) return np.clip(np.squeeze(np.stack(chls, -1)), -1.0, 1.0) def decode_waveform(self, lat, dec, dec2, batch_size=64): lat = lat[:, :, : (lat.shape[-2] // 8) * 8, :] abls = tf.split(lat, lat.shape[-2] // 8, -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=batch_size, ) 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=batch_size) abwv = self.conc_tog_specphase(ab_m, ab_p) chls.append(abwv) return np.clip(np.squeeze(np.stack(chls, -1)), -1.0, 1.0) def get_noise_interp_multi(self, fac=1, var=2.0): noiseg = self.truncated_normal([1, self.args.coorddepth], var, dtype=tf.float32) coordratio = self.args.coordlen // self.args.latlen noisels = [ tf.concat([self.truncated_normal([1, 64], var, dtype=tf.float32), noiseg], -1) for i in range(3 + ((fac - 1) // coordratio)) ] rls = tf.concat( [ tf.linspace(noisels[k], noisels[k + 1], self.args.coordlen + 1, axis=-2)[:, :-1, :] for k in range(len(noisels) - 1) ], -2, ) rls = self.center_coordinate(rls) rls = rls[:, self.args.latlen // 4 :, :] rls = rls[:, : (rls.shape[-2] // self.args.latlen) * self.args.latlen, :] rls = tf.split(rls, rls.shape[-2] // self.args.latlen, -2) return tf.concat(rls[:fac], 0) def get_noise_interp_loop(self, fac=1, var=2.0): noiseg = self.truncated_normal([1, self.args.coorddepth], var, dtype=tf.float32) coordratio = self.args.coordlen // self.args.latlen noisels_pre = [tf.concat([self.truncated_normal([1, 64], var, dtype=tf.float32), noiseg], -1) for i in range(2)] noisels = [] for k in range(fac + 2): noisels.append(noisels_pre[0]) noisels.append(noisels_pre[1]) rls = tf.concat( [ tf.linspace(noisels[k], noisels[k + 1], self.args.latlen // 2 + 1, axis=-2)[:, :-1, :] for k in range(len(noisels) - 1) ], -2, ) rls = self.center_coordinate(rls) rls = rls[:, self.args.latlen // 2 :, :] rls = rls[:, : (rls.shape[-2] // self.args.latlen) * self.args.latlen, :] rls = tf.split(rls, rls.shape[-2] // self.args.latlen, -2) return tf.concat(rls[:fac], 0) def generate(self, models_ls): critic, gen, enc, dec, enc2, dec2, gen_ema, [opt_dec, opt_disc], switch = models_ls os.makedirs(self.args.save_path, exist_ok=True) fac = (self.args.seconds // 23) + 1 print(f"Generating {self.args.num_samples} samples...") for i in tqdm(range(self.args.num_samples)): wv = self.generate_waveform( self.get_noise_interp_multi(fac, self.args.truncation), gen_ema, dec, dec2, batch_size=64 ) dt = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") write_wav( f"{self.args.save_path}/{i}_{dt}.wav", self.args.sr, np.squeeze(wv)[: self.args.seconds * self.args.sr] ) def decode_path(self, models_ls): critic, gen, enc, dec, enc2, dec2, gen_ema, [opt_dec, opt_disc], switch = models_ls os.makedirs(self.args.save_path, exist_ok=True) pathls = glob(self.args.files_path + "/*.npy") print(f"Decoding {len(pathls)} samples...") for p in tqdm(pathls): tp, ext = os.path.splitext(p) bname = os.path.basename(tp) lat = np.load(p, allow_pickle=True) lat = tf.expand_dims(lat, 0) lat = tf.expand_dims(lat, 0) wv = self.decode_waveform(lat, dec, dec2, batch_size=64) # dt = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") write_wav(f"{self.args.save_path}/{bname}.wav", self.args.sr, np.squeeze(wv)) def stfunc(self, genre, z, var, models_ls_1, models_ls_2, models_ls_3): critic, gen, enc, dec, enc2, dec2, gen_ema_1, [opt_dec, opt_disc], switch = models_ls_1 critic, gen, enc, dec, enc2, dec2, gen_ema_2, [opt_dec, opt_disc], switch = models_ls_2 critic, gen, enc, dec, enc2, dec2, gen_ema_3, [opt_dec, opt_disc], switch = models_ls_3 if genre == 0: gen_ema = gen_ema_1 elif genre == 1: gen_ema = gen_ema_2 else: gen_ema = gen_ema_3 var = float(var) if z == 0: fac = 1 elif z == 1: fac = 5 else: fac = 10 bef = time.time() noiseinp = self.get_noise_interp_multi(fac, var) abwvc = self.generate_waveform(noiseinp, gen_ema, dec, dec2, batch_size=64) # print( # f"Time for complete generation pipeline: {time.time()-bef} s {int(np.round((fac*23.)/(time.time()-bef)))}x faster than Real Time!" # ) 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, self.args.hop * 2 ), [1, 0], ) ), -2, ) #output = "/tmp/outputfile.wav" #write_wav (output,self.args.sr,np.int16(abwvc * 32767.0)) return ( (self.args.sr, np.int16(abwvc * 32767.0)) ) def render_gradio(self, models_ls_1, models_ls_2, models_ls_3, train=True): article_text = "Original work by Marco Pasini ([Twitter](https://twitter.com/marco_ppasini)) at the Institute of Computational Perception, JKU Linz. Supervised by Jan Schlüter." def gradio_func(genre, x, y): return self.stfunc(genre, x, y, models_ls_1, models_ls_2, models_ls_3) if self.args.small: durations = ["11s", "59s", "1m 58s"] durations_default = "59s" else: durations = ["23s", "1m 58s", "3m 57s"] durations_default = "1m 58s" iface = gr.Interface( fn=gradio_func, inputs=[ gr.Radio( choices=["Techno/Experimental", "Nes Acoustic Energy (finetuned)", "Misc"], type="index", value="Techno/Experimental", label="Music Genre to Generate", ), gr.Radio( choices=durations, type="index", value=durations_default, label="Generated Music Length", ), gr.Slider( minimum=0.1, maximum=3.9, step=0.1, value=1.8, label="How much do you want the music style to be varied? (Stddev truncation for random vectors)", ), ], outputs=[ gr.Audio(label="output audio", type="numpy") ], title="musika!", description="Forked from https://huggingface.co/spaces/marcop/musika I could not get the API to work on the original version so I have this one just exporting audio files. Should be adding some models next. Blazingly Fast 44.1 kHz Stereo Waveform Music Generation of Arbitrary Length. Be patient and enjoy the weirdness!", article=article_text, ) 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("--------------------------------")