Add first version

README.md

@@ -4,7 +4,7 @@ emoji: 🚀
 colorFrom: purple
 colorTo: blue
 sdk: gradio
-sdk_version: 3.0.24
+sdk_version: 2.5.2
 app_file: app.py
 pinned: false
 license: cc-by-4.0

app.py

@@ -0,0 +1,15 @@
+from parse_test import parse_args
+from models import Models_functions
+from utils import Utils_functions
+
+
+# parse args
+args = parse_args()
+
+# initialize networks
+M = Models_functions(args)
+models_ls_techno, models_ls_classical = M.get_networks()
+
+# test musika
+U = Utils_functions(args)
+U.render_gradio(models_ls_techno, models_ls_classical, train=False)

checkpoints/classical/dec.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32e50f938685f56e1dc5e137e95c59472418194125e435cafb668584e65b0fcc
+size 29745512

checkpoints/classical/dec2.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:643e1415f5ebf34cb067b2ccaced657f1f3e3f068810c2d130dfc3abb20c3cc2
+size 26799136

checkpoints/classical/gen_ema.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1508fa1fad225f7c29ed64252b531e4462c97b0f0c80dd6d212014916f6261eb
+size 56431944

checkpoints/techno/dec.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:953f70f8c874d36fd57dfd825b3bceff2f7dba04c93f4a1464c0c5d223c2b2e7
+size 29745512

checkpoints/techno/dec2.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd0a9e2eb78bae5178dd25022b4d0ac5020a6366c98fa009830c6c8818582045
+size 26799136

checkpoints/techno/gen_ema.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9a81012a6bc318e23bb12ed0fc70260d260e71c9b355b553593c60bfec0fab41
+size 56431944

layers.py

@@ -0,0 +1,164 @@
+import tensorflow as tf
+import tensorflow.keras.backend as K
+from tensorflow.keras.layers import Conv2D, Conv2DTranspose, Dense
+from tensorflow.python.eager import context
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.keras import activations, constraints, initializers, regularizers
+from tensorflow.python.keras.layers.convolutional import SeparableConv
+from tensorflow.python.ops import (
+    array_ops,
+    gen_math_ops,
+    math_ops,
+    sparse_ops,
+    standard_ops,
+)
+
+
+def l2normalize(v, eps=1e-12):
+    return v / (tf.norm(v) + eps)
+
+
+class ConvSN2D(tf.keras.layers.Conv2D):
+    def __init__(self, filters, kernel_size, power_iterations=1, **kwargs):
+        super(ConvSN2D, self).__init__(filters, kernel_size, **kwargs)
+        self.power_iterations = power_iterations
+
+    def build(self, input_shape):
+        super(ConvSN2D, self).build(input_shape)
+
+        if self.data_format == "channels_first":
+            channel_axis = 1
+        else:
+            channel_axis = -1
+
+        self.u = self.add_weight(
+            self.name + "_u",
+            shape=tuple([1, self.kernel.shape.as_list()[-1]]),
+            initializer=tf.initializers.RandomNormal(0, 1),
+            trainable=False,
+            dtype=self.dtype,
+        )
+
+    def compute_spectral_norm(self, W, new_u, W_shape):
+        for _ in range(self.power_iterations):
+
+            new_v = l2normalize(tf.matmul(new_u, tf.transpose(W)))
+            new_u = l2normalize(tf.matmul(new_v, W))
+
+        sigma = tf.matmul(tf.matmul(new_v, W), tf.transpose(new_u))
+        W_bar = W / sigma
+
+        with tf.control_dependencies([self.u.assign(new_u)]):
+            W_bar = tf.reshape(W_bar, W_shape)
+
+        return W_bar
+
+    def call(self, inputs):
+        W_shape = self.kernel.shape.as_list()
+        W_reshaped = tf.reshape(self.kernel, (-1, W_shape[-1]))
+        new_kernel = self.compute_spectral_norm(W_reshaped, self.u, W_shape)
+        outputs = self._convolution_op(inputs, new_kernel)
+
+        if self.use_bias:
+            if self.data_format == "channels_first":
+                outputs = tf.nn.bias_add(outputs, self.bias, data_format="NCHW")
+            else:
+                outputs = tf.nn.bias_add(outputs, self.bias, data_format="NHWC")
+        if self.activation is not None:
+            return self.activation(outputs)
+
+        return outputs
+
+
+class DenseSN(Dense):
+    def build(self, input_shape):
+        super(DenseSN, self).build(input_shape)
+
+        self.u = self.add_weight(
+            self.name + "_u",
+            shape=tuple([1, self.kernel.shape.as_list()[-1]]),
+            initializer=tf.initializers.RandomNormal(0, 1),
+            trainable=False,
+            dtype=self.dtype,
+        )
+
+    def compute_spectral_norm(self, W, new_u, W_shape):
+        new_v = l2normalize(tf.matmul(new_u, tf.transpose(W)))
+        new_u = l2normalize(tf.matmul(new_v, W))
+        sigma = tf.matmul(tf.matmul(new_v, W), tf.transpose(new_u))
+        W_bar = W / sigma
+        with tf.control_dependencies([self.u.assign(new_u)]):
+            W_bar = tf.reshape(W_bar, W_shape)
+        return W_bar
+
+    def call(self, inputs):
+        W_shape = self.kernel.shape.as_list()
+        W_reshaped = tf.reshape(self.kernel, (-1, W_shape[-1]))
+        new_kernel = self.compute_spectral_norm(W_reshaped, self.u, W_shape)
+        rank = len(inputs.shape)
+        if rank > 2:
+            outputs = standard_ops.tensordot(inputs, new_kernel, [[rank - 1], [0]])
+            if not context.executing_eagerly():
+                shape = inputs.shape.as_list()
+                output_shape = shape[:-1] + [self.units]
+                outputs.set_shape(output_shape)
+        else:
+            inputs = math_ops.cast(inputs, self._compute_dtype)
+            if K.is_sparse(inputs):
+                outputs = sparse_ops.sparse_tensor_dense_matmul(inputs, new_kernel)
+            else:
+                outputs = gen_math_ops.mat_mul(inputs, new_kernel)
+        if self.use_bias:
+            outputs = tf.nn.bias_add(outputs, self.bias)
+        if self.activation is not None:
+            return self.activation(outputs)
+        return outputs
+
+
+class AddNoise(tf.keras.layers.Layer):
+    def build(self, input_shape):
+        self.b = self.add_weight(
+            shape=[
+                1,
+            ],
+            initializer=tf.keras.initializers.zeros(),
+            trainable=True,
+            name="noise_weight",
+        )
+
+    def call(self, inputs):
+        rand = tf.random.normal(
+            [tf.shape(inputs)[0], inputs.shape[1], inputs.shape[2], 1],
+            mean=0.0,
+            stddev=1.0,
+            dtype=self.dtype,
+        )
+        output = inputs + self.b * rand
+        return output
+
+
+class PosEnc(tf.keras.layers.Layer):
+    def __init__(self, **kwargs):
+        super(PosEnc, self).__init__(**kwargs)
+
+    def call(self, inputs):
+        # inputs shape: [bs,mel_bins,shape,1]
+        pos = tf.repeat(
+            tf.reshape(tf.range(inputs.shape[-3], dtype=tf.int32), [1, -1, 1, 1]),
+            inputs.shape[-2],
+            -2,
+        )
+        pos = tf.cast(tf.repeat(pos, tf.shape(inputs)[0], 0), self.dtype) / tf.cast(
+            inputs.shape[-3], self.dtype
+        )
+        return tf.concat([inputs, pos], -1)  # [bs,1,hop,2]
+
+
+def flatten_hw(x, data_format="channels_last"):
+    if data_format == "channels_last":
+        x = tf.transpose(x, perm=[0, 3, 1, 2])  # Convert to `channels_first`
+
+    old_shape = tf.shape(x)
+    new_shape = [old_shape[0], old_shape[2] * old_shape[3], old_shape[1]]
+
+    return tf.reshape(x, new_shape)

models.py

@@ -0,0 +1,831 @@
+import numpy as np
+import tensorflow as tf
+import tensorflow_addons as tfa
+from tensorflow.keras import mixed_precision
+from tensorflow.keras.layers import (
+    Add,
+    BatchNormalization,
+    Concatenate,
+    Conv2D,
+    Conv2DTranspose,
+    Cropping1D,
+    Cropping2D,
+    Dense,
+    Dot,
+    Flatten,
+    GlobalAveragePooling2D,
+    Input,
+    Lambda,
+    LeakyReLU,
+    Multiply,
+    ReLU,
+    Reshape,
+    SeparableConv2D,
+    UpSampling2D,
+    ZeroPadding2D,
+)
+from tensorflow.keras.models import Model, Sequential
+from tensorflow.keras.optimizers import Adam
+from tensorflow.python.keras.utils.layer_utils import count_params
+
+from layers import ConvSN2D, DenseSN, PosEnc, AddNoise
+
+
+class Models_functions:
+    def __init__(self, args):
+
+        self.args = args
+        if self.args.mixed_precision:
+            self.mixed_precision = mixed_precision
+            self.policy = mixed_precision.Policy("mixed_float16")
+            mixed_precision.set_global_policy(self.policy)
+        self.init = tf.keras.initializers.he_uniform()
+
+    def conv_util(
+        self,
+        inp,
+        filters,
+        kernel_size=(1, 3),
+        strides=(1, 1),
+        noise=False,
+        upsample=False,
+        padding="same",
+        bn=True,
+    ):
+
+        x = inp
+
+        if upsample:
+            x = tf.keras.layers.Conv2DTranspose(
+                filters,
+                kernel_size=kernel_size,
+                strides=strides,
+                activation="linear",
+                padding=padding,
+                kernel_initializer=self.init,
+            )(x)
+        else:
+            x = tf.keras.layers.Conv2D(
+                filters,
+                kernel_size=kernel_size,
+                strides=strides,
+                activation="linear",
+                padding=padding,
+                kernel_initializer=self.init,
+            )(x)
+
+        if noise:
+            x = AddNoise()(x)
+
+        if bn:
+            x = tf.keras.layers.BatchNormalization()(x)
+
+        x = tf.keras.activations.swish(x)
+
+        return x
+
+    def adain(self, x, emb):
+        emb = tf.keras.layers.Conv2D(
+            x.shape[-1],
+            kernel_size=(1, 1),
+            strides=1,
+            activation="linear",
+            padding="same",
+            kernel_initializer=self.init,
+            use_bias=True,
+        )(emb)
+        x = x / (tf.math.reduce_std(x, -2, keepdims=True) + 1e-7)
+        return x * emb
+
+    def se_layer(self, x, filters):
+        x = tf.reduce_mean(x, -2, keepdims=True)
+        x = tf.keras.layers.Conv2D(
+            filters,
+            kernel_size=(1, 1),
+            strides=(1, 1),
+            activation="linear",
+            padding="valid",
+            kernel_initializer=self.init,
+            use_bias=True,
+        )(x)
+        x = tf.keras.activations.swish(x)
+        return tf.keras.layers.Conv2D(
+            filters,
+            kernel_size=(1, 1),
+            strides=(1, 1),
+            activation="sigmoid",
+            padding="valid",
+            kernel_initializer=self.init,
+            use_bias=True,
+        )(x)
+
+    def conv_util_gen(
+        self,
+        inp,
+        filters,
+        kernel_size=(1, 9),
+        strides=(1, 1),
+        noise=False,
+        upsample=False,
+        emb=None,
+        se1=None,
+    ):
+
+        x = inp
+
+        if upsample:
+            x = tf.keras.layers.Conv2DTranspose(
+                filters,
+                kernel_size=kernel_size,
+                strides=strides,
+                activation="linear",
+                padding="same",
+                kernel_initializer=self.init,
+                use_bias=True,
+            )(x)
+        else:
+            x = tf.keras.layers.Conv2D(
+                filters,
+                kernel_size=kernel_size,
+                strides=strides,
+                activation="linear",
+                padding="same",
+                kernel_initializer=self.init,
+                use_bias=True,
+            )(x)
+
+        if noise:
+            x = AddNoise()(x)
+
+        if emb is not None:
+            x = self.adain(x, emb)
+        else:
+            x = tf.keras.layers.BatchNormalization()(x)
+
+        x1 = tf.keras.activations.swish(x)
+
+        if se1 is not None:
+            x1 = x1 * se1
+
+        return x1
+
+    def res_block_disc(self, inp, filters, kernel_size=(1, 3), kernel_size_2=None, strides=(1, 1)):
+
+        if kernel_size_2 is None:
+            kernel_size_2 = kernel_size
+
+        x = tf.keras.layers.Conv2D(
+            inp.shape[-1],
+            kernel_size=kernel_size_2,
+            strides=1,
+            activation="linear",
+            padding="same",
+            kernel_initializer=self.init,
+        )(inp)
+        x = tf.keras.layers.LeakyReLU(0.2)(x)
+        x = tf.math.sqrt(tf.cast(0.5, self.args.datatype)) * x
+        x = tf.keras.layers.Conv2D(
+            filters,
+            kernel_size=kernel_size,
+            strides=strides,
+            activation="linear",
+            padding="same",
+            kernel_initializer=self.init,
+        )(x)
+        x = tf.keras.layers.LeakyReLU(0.2)(x)
+        x = tf.math.sqrt(tf.cast(0.5, self.args.datatype)) * x
+
+        if strides != (1, 1):
+            inp = tf.keras.layers.AveragePooling2D(strides, padding="same")(inp)
+
+        if inp.shape[-1] != filters:
+            inp = tf.keras.layers.Conv2D(
+                filters,
+                kernel_size=1,
+                strides=1,
+                activation="linear",
+                padding="same",
+                kernel_initializer=self.init,
+                use_bias=False,
+            )(inp)
+
+        return x + inp
+
+    def build_encoder2(self):
+
+        dim = 128
+
+        inpf = Input((1, self.args.shape, dim))
+
+        inpfls = tf.split(inpf, 16, -2)
+        inpb = tf.concat(inpfls, 0)
+
+        g0 = self.conv_util(inpb, 256, kernel_size=(1, 1), strides=(1, 1), padding="valid")
+        g1 = self.conv_util(g0, 256 + 256, kernel_size=(1, 3), strides=(1, 1), padding="valid")
+        g2 = self.conv_util(g1, 512 + 128, kernel_size=(1, 3), strides=(1, 1), padding="valid")
+        g3 = self.conv_util(g2, 512 + 128, kernel_size=(1, 1), strides=(1, 1), padding="valid")
+        g4 = self.conv_util(g3, 512 + 256, kernel_size=(1, 3), strides=(1, 1), padding="valid")
+        g5 = self.conv_util(g4, 512 + 256, kernel_size=(1, 2), strides=(1, 1), padding="valid")
+
+        g = tf.keras.layers.Conv2D(
+            64,
+            kernel_size=(1, 1),
+            strides=1,
+            padding="valid",
+            kernel_initializer=self.init,
+            name="cbottle",
+            activation="tanh",
+        )(g5)
+
+        gls = tf.split(g, 16, 0)
+        g = tf.concat(gls, -2)
+        gls = tf.split(g, 2, -2)
+        g = tf.concat(gls, 0)
+
+        gf = tf.cast(g, tf.float32)
+        return Model(inpf, gf, name="ENC2")
+
+    def build_decoder2(self):
+
+        dim = 128
+        bottledim = 64
+
+        inpf = Input((1, self.args.shape // 16, bottledim))
+
+        g = inpf
+
+        g = self.conv_util(
+            g,
+            512 + 128 + 128,
+            kernel_size=(1, 4),
+            strides=(1, 1),
+            upsample=False,
+            noise=True,
+        )
+        g = self.conv_util(
+            g,
+            512 + 128 + 128,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            upsample=True,
+            noise=True,
+        )
+        g = self.conv_util(g, 512 + 128, kernel_size=(1, 4), strides=(1, 2), upsample=True, noise=True)
+        g = self.conv_util(g, 512, kernel_size=(1, 4), strides=(1, 1), upsample=False, noise=True)
+        g = self.conv_util(g, 256 + 128, kernel_size=(1, 4), strides=(1, 2), upsample=True, noise=True)
+
+        gf = tf.keras.layers.Conv2D(
+            dim,
+            kernel_size=(1, 1),
+            strides=1,
+            padding="same",
+            activation="tanh",
+            kernel_initializer=self.init,
+        )(g)
+
+        gfls = tf.split(gf, 2, 0)
+        gf = tf.concat(gfls, -2)
+
+        gf = tf.cast(gf, tf.float32)
+
+        return Model(inpf, gf, name="DEC2")
+
+    def build_encoder(self):
+
+        dim = ((4 * self.args.hop) // 2) + 1
+
+        inpf = Input((dim, self.args.shape, 1))
+
+        ginp = tf.transpose(inpf, [0, 3, 2, 1])
+
+        g0 = self.conv_util(
+            ginp,
+            self.args.hop * 2 + 32,
+            kernel_size=(1, 1),
+            strides=(1, 1),
+            padding="valid",
+        )
+
+        g = self.conv_util(
+            g0,
+            self.args.hop * 2 + 64,
+            kernel_size=(1, 1),
+            strides=(1, 1),
+            padding="valid",
+        )
+        g = self.conv_util(
+            g,
+            self.args.hop * 2 + 64 + 64,
+            kernel_size=(1, 1),
+            strides=(1, 1),
+            padding="valid",
+        )
+        g = self.conv_util(
+            g,
+            self.args.hop * 2 + 128 + 64,
+            kernel_size=(1, 1),
+            strides=(1, 1),
+            padding="valid",
+        )
+        g = self.conv_util(
+            g,
+            self.args.hop * 2 + 128 + 128,
+            kernel_size=(1, 1),
+            strides=(1, 1),
+            padding="valid",
+        )
+
+        g = tf.keras.layers.Conv2D(
+            128,
+            kernel_size=(1, 1),
+            strides=1,
+            padding="valid",
+            kernel_initializer=self.init,
+        )(g)
+        gb = tf.keras.activations.tanh(g)
+
+        gbls = tf.split(gb, 2, -2)
+        gb = tf.concat(gbls, 0)
+
+        gb = tf.cast(gb, tf.float32)
+        return Model(inpf, gb, name="ENC")
+
+    def build_decoder(self):
+
+        dim = ((4 * self.args.hop) // 2) + 1
+
+        inpf = Input((1, self.args.shape // 2, 128))
+
+        g = inpf
+
+        g0 = self.conv_util(g, self.args.hop * 3, kernel_size=(1, 1), strides=(1, 1), noise=True)
+
+        g1 = self.conv_util(g0, self.args.hop * 2, kernel_size=(1, 3), strides=(1, 2), noise=True)
+        g2 = self.conv_util(
+            g1,
+            self.args.hop + self.args.hop // 2,
+            kernel_size=(1, 3),
+            strides=(1, 2),
+            noise=True,
+        )
+        g = self.conv_util(
+            g2,
+            self.args.hop + self.args.hop // 4,
+            kernel_size=(1, 3),
+            strides=(1, 2),
+            noise=True,
+        )
+
+        g = self.conv_util(
+            g,
+            self.args.hop + self.args.hop // 2,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            upsample=True,
+            noise=True,
+        )
+        g = self.conv_util(
+            g + g2,
+            self.args.hop * 2,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            upsample=True,
+            noise=True,
+        )
+        g = self.conv_util(
+            g + g1,
+            self.args.hop * 3,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            upsample=True,
+            noise=True,
+        )
+
+        g = self.conv_util(g + g0, self.args.hop * 5, kernel_size=(1, 1), strides=(1, 1), noise=True)
+
+        g = Conv2D(
+            dim * 2,
+            kernel_size=(1, 1),
+            strides=(1, 1),
+            kernel_initializer=self.init,
+            padding="same",
+        )(g)
+        g = tf.clip_by_value(g, -1.0, 1.0)
+
+        gf, pf = tf.split(g, 2, -1)
+
+        gfls = tf.split(gf, self.args.shape // self.args.window, 0)
+        gf = tf.concat(gfls, -2)
+
+        pfls = tf.split(pf, self.args.shape // self.args.window, 0)
+        pf = tf.concat(pfls, -2)
+
+        s = tf.transpose(gf, [0, 2, 3, 1])
+        p = tf.transpose(pf, [0, 2, 3, 1])
+
+        s = tf.cast(tf.squeeze(s, -1), tf.float32)
+        p = tf.cast(tf.squeeze(p, -1), tf.float32)
+
+        return Model(inpf, [s, p], name="DEC")
+
+    def build_critic(self):
+
+        sinp = Input(shape=(1, self.args.latlen, self.args.latdepth * 2))
+
+        dim = 64 * 2
+
+        sf = tf.keras.layers.Conv2D(
+            self.args.latdepth * 4,
+            kernel_size=(1, 1),
+            strides=(1, 1),
+            activation="linear",
+            padding="valid",
+            kernel_initializer=self.init,
+            use_bias=False,
+            trainable=False,
+        )(sinp)
+
+        sf = tf.keras.layers.Conv2D(
+            256 + 128,
+            kernel_size=(1, 3),
+            strides=(1, 2),
+            activation="linear",
+            padding="same",
+            kernel_initializer=self.init,
+        )(sf)
+        sf = tf.keras.layers.LeakyReLU(0.2)(sf)
+        sf = self.res_block_disc(sf, 256 + 128 + 128, kernel_size=(1, 3), strides=(1, 2))
+        sf = self.res_block_disc(sf, 512 + 128, kernel_size=(1, 3), strides=(1, 2))
+        sf = self.res_block_disc(sf, 512 + 256, kernel_size=(1, 3), strides=(1, 2))
+        sf = self.res_block_disc(sf, 512 + 128 + 256, kernel_size=(1, 3), strides=(1, 2))
+        sfo = self.res_block_disc(sf, 512 + 512, kernel_size=(1, 3), strides=(1, 2), kernel_size_2=(1, 1))
+        sf = sfo
+
+        gf = tf.keras.layers.Dense(1, activation="linear", use_bias=True, kernel_initializer=self.init)(Flatten()(sf))
+
+        gf = tf.cast(gf, tf.float32)
+        sfo = tf.cast(sfo, tf.float32)
+
+        return Model(sinp, [gf, sfo], name="C")
+
+    def build_critic_rec(self):
+
+        sinp = Input(shape=(1, self.args.latlen // 64, 512 + 512))
+
+        dim = self.args.latdepth * 2
+
+        sf = tf.keras.layers.Conv2DTranspose(
+            512,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            activation="linear",
+            padding="same",
+            kernel_initializer=self.init,
+        )(sinp)
+        sf = tf.keras.layers.LeakyReLU(0.2)(sf)
+
+        sf = tf.keras.layers.Conv2DTranspose(
+            256 + 128 + 64,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            activation="linear",
+            padding="same",
+            kernel_initializer=self.init,
+        )(sf)
+        sf = tf.keras.layers.LeakyReLU(0.2)(sf)
+        sf = tf.keras.layers.Conv2DTranspose(
+            256 + 128,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            activation="linear",
+            padding="same",
+            kernel_initializer=self.init,
+        )(sf)
+        sf = tf.keras.layers.LeakyReLU(0.2)(sf)
+        sf = tf.keras.layers.Conv2DTranspose(
+            256 + 64,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            activation="linear",
+            padding="same",
+            kernel_initializer=self.init,
+        )(sf)
+        sf = tf.keras.layers.LeakyReLU(0.2)(sf)
+        sf = tf.keras.layers.Conv2DTranspose(
+            256,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            activation="linear",
+            padding="same",
+            kernel_initializer=self.init,
+        )(sf)
+        sf = tf.keras.layers.LeakyReLU(0.2)(sf)
+        sf = tf.keras.layers.Conv2DTranspose(
+            128 + 64,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            activation="linear",
+            padding="same",
+            kernel_initializer=self.init,
+        )(sf)
+        sf = tf.keras.layers.LeakyReLU(0.2)(sf)
+
+        gf = tf.keras.layers.Conv2D(
+            dim,
+            kernel_size=(1, 1),
+            strides=(1, 1),
+            activation="tanh",
+            padding="same",
+            kernel_initializer=self.init,
+        )(sf)
+
+        gf = tf.cast(gf, tf.float32)
+
+        return Model(sinp, gf, name="CR")
+
+    def build_generator(self):
+
+        dim = self.args.latdepth * 2
+
+        inpf = Input((self.args.latlen, self.args.latdepth * 2))
+
+        inpfls = tf.split(inpf, 2, -2)
+        inpb = tf.concat(inpfls, 0)
+
+        inpg = tf.reduce_mean(inpb, -2)
+        inp1 = tf.keras.layers.AveragePooling2D((1, 2), padding="valid")(tf.expand_dims(inpb, -3))
+        inp2 = tf.keras.layers.AveragePooling2D((1, 2), padding="valid")(inp1)
+        inp3 = tf.keras.layers.AveragePooling2D((1, 2), padding="valid")(inp2)
+        inp4 = tf.keras.layers.AveragePooling2D((1, 2), padding="valid")(inp3)
+        inp5 = tf.keras.layers.AveragePooling2D((1, 2), padding="valid")(inp4)
+        inp6 = tf.keras.layers.AveragePooling2D((1, 2), padding="valid")(inp5)
+
+        g = tf.keras.layers.Dense(
+            4 * (512 + 256 + 128),
+            activation="linear",
+            use_bias=True,
+            kernel_initializer=self.init,
+        )(Flatten()(inp6))
+        g = tf.keras.layers.Reshape((1, 4, 512 + 256 + 128))(g)
+        g = AddNoise()(g)
+        g = self.adain(g, inp5)
+        g = tf.keras.activations.swish(g)
+
+        g = self.conv_util_gen(
+            g,
+            512 + 256,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            upsample=True,
+            noise=True,
+            emb=inp4,
+        )
+        g1 = self.conv_util_gen(
+            g,
+            512 + 256,
+            kernel_size=(1, 1),
+            strides=(1, 1),
+            upsample=False,
+            noise=True,
+            emb=inp4,
+        )
+        g2 = self.conv_util_gen(
+            g1,
+            512 + 128,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            upsample=True,
+            noise=True,
+            emb=inp3,
+        )
+        g2b = self.conv_util_gen(
+            g2,
+            512 + 128,
+            kernel_size=(1, 3),
+            strides=(1, 1),
+            upsample=False,
+            noise=True,
+            emb=inp3,
+        )
+        g3 = self.conv_util_gen(
+            g2b,
+            256 + 256,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            upsample=True,
+            noise=True,
+            emb=inp2,
+            se1=self.se_layer(g, 256 + 256),
+        )
+        g3 = self.conv_util_gen(
+            g3,
+            256 + 256,
+            kernel_size=(1, 3),
+            strides=(1, 1),
+            upsample=False,
+            noise=True,
+            emb=inp2,
+            se1=self.se_layer(g1, 256 + 256),
+        )
+        g4 = self.conv_util_gen(
+            g3,
+            256 + 128,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            upsample=True,
+            noise=True,
+            emb=inp1,
+            se1=self.se_layer(g2, 256 + 128),
+        )
+        g4 = self.conv_util_gen(
+            g4,
+            256 + 128,
+            kernel_size=(1, 3),
+            strides=(1, 1),
+            upsample=False,
+            noise=True,
+            emb=inp1,
+            se1=self.se_layer(g2b, 256 + 128),
+        )
+        g5 = self.conv_util_gen(
+            g4,
+            256,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            upsample=True,
+            noise=True,
+            emb=tf.expand_dims(tf.cast(inpb, dtype=self.args.datatype), -3),
+        )
+
+        gf = tf.keras.layers.Conv2D(
+            dim,
+            kernel_size=(1, 1),
+            strides=(1, 1),
+            kernel_initializer=self.init,
+            padding="same",
+            activation="tanh",
+        )(g5)
+
+        gfls = tf.split(gf, 2, 0)
+        gf = tf.concat(gfls, -2)
+
+        gf = tf.cast(gf, tf.float32)
+
+        return Model(inpf, gf, name="GEN")
+
+    # Load past models from path to resume training or test
+    def load(self, path, load_dec=False):
+        gen = self.build_generator()
+        critic = self.build_critic()
+        enc = self.build_encoder()
+        dec = self.build_decoder()
+        enc2 = self.build_encoder2()
+        dec2 = self.build_decoder2()
+        critic_rec = self.build_critic_rec()
+        gen_ema = self.build_generator()
+
+        if self.args.mixed_precision:
+            opt_disc = self.mixed_precision.LossScaleOptimizer(tf.keras.optimizers.Adam(0.0001, 0.9))
+            opt_dec = self.mixed_precision.LossScaleOptimizer(tf.keras.optimizers.Adam(0.0001, 0.9))
+        else:
+            opt_disc = tf.keras.optimizers.Adam(0.0001, 0.9)
+            opt_dec = tf.keras.optimizers.Adam(0.0001, 0.9)
+
+        if load_dec:
+            dec.load_weights(self.args.dec_path + "/dec.h5")
+            dec2.load_weights(self.args.dec_path + "/dec2.h5")
+
+        else:
+            grad_vars = critic.trainable_weights + critic_rec.trainable_weights
+            zero_grads = [tf.zeros_like(w) for w in grad_vars]
+            opt_disc.apply_gradients(zip(zero_grads, grad_vars))
+
+            grad_vars = gen.trainable_variables
+            zero_grads = [tf.zeros_like(w) for w in grad_vars]
+            opt_dec.apply_gradients(zip(zero_grads, grad_vars))
+
+            if not self.args.testing:
+                opt_disc.set_weights(np.load(path + "/opt_disc.npy", allow_pickle=True))
+                opt_dec.set_weights(np.load(path + "/opt_dec.npy", allow_pickle=True))
+
+            if not self.args.testing:
+                critic.load_weights(path + "/critic.h5")
+                gen.load_weights(path + "/gen.h5")
+                enc.load_weights(path + "/enc.h5")
+                enc2.load_weights(path + "/enc2.h5")
+                critic_rec.load_weights(path + "/critic_rec.h5")
+            gen_ema.load_weights(path + "/gen_ema.h5")
+            dec.load_weights(path + "/dec.h5")
+            dec2.load_weights(path + "/dec2.h5")
+
+        return (
+            critic,
+            gen,
+            enc,
+            dec,
+            enc2,
+            dec2,
+            critic_rec,
+            gen_ema,
+            [opt_dec, opt_disc],
+        )
+
+    def build(self):
+        gen = self.build_generator()
+        critic = self.build_critic()
+        enc = self.build_encoder()
+        dec = self.build_decoder()
+        enc2 = self.build_encoder2()
+        dec2 = self.build_decoder2()
+        critic_rec = self.build_critic_rec()
+        gen_ema = self.build_generator()
+
+        gen_ema = tf.keras.models.clone_model(gen)
+        gen_ema.set_weights(gen.get_weights())
+
+        if self.args.mixed_precision:
+            opt_disc = self.mixed_precision.LossScaleOptimizer(tf.keras.optimizers.Adam(0.0001, 0.9))
+            opt_dec = self.mixed_precision.LossScaleOptimizer(tf.keras.optimizers.Adam(0.0001, 0.9))
+        else:
+            opt_disc = tf.keras.optimizers.Adam(0.0001, 0.9)
+            opt_dec = tf.keras.optimizers.Adam(0.0001, 0.9)
+
+        return (
+            critic,
+            gen,
+            enc,
+            dec,
+            enc2,
+            dec2,
+            critic_rec,
+            gen_ema,
+            [opt_dec, opt_disc],
+        )
+
+    def get_networks(self):
+        (
+            critic,
+            gen,
+            enc,
+            dec_techno,
+            enc2,
+            dec2_techno,
+            critic_rec,
+            gen_ema_techno,
+            [opt_dec, opt_disc],
+        ) = self.load(self.args.load_path_techno, load_dec=False)
+        print(f"Techno networks loaded from {self.args.load_path_techno}")
+
+        (
+            critic,
+            gen,
+            enc,
+            dec_classical,
+            enc2,
+            dec2_classical,
+            critic_rec,
+            gen_ema_classical,
+            [opt_dec, opt_disc],
+        ) = self.load(self.args.load_path_classical, load_dec=False)
+        print(f"Classical networks loaded from {self.args.load_path_classical}")
+
+        return [critic, gen, enc, dec_techno, enc2, dec2_techno, critic_rec, gen_ema_techno, [opt_dec, opt_disc]], [
+            critic,
+            gen,
+            enc,
+            dec_classical,
+            enc2,
+            dec2_classical,
+            critic_rec,
+            gen_ema_classical,
+            [opt_dec, opt_disc],
+        ]
+
+    def initialize_networks(self):
+
+        [critic, gen, enc, dec_techno, enc2, dec2_techno, critic_rec, gen_ema_techno, [opt_dec, opt_disc]], [
+            critic,
+            gen,
+            enc,
+            dec_classical,
+            enc2,
+            dec2_classical,
+            critic_rec,
+            gen_ema_classical,
+            [opt_dec, opt_disc],
+        ] = self.get_networks()
+
+        print(f"Generator params: {count_params(gen_ema_techno.trainable_variables)}")
+        print(f"Decoder params: {count_params(dec_techno.trainable_variables+dec2_techno.trainable_variables)}")
+
+        return [critic, gen, enc, dec_techno, enc2, dec2_techno, critic_rec, gen_ema_techno, [opt_dec, opt_disc]], [
+            critic,
+            gen,
+            enc,
+            dec_classical,
+            enc2,
+            dec2_classical,
+            critic_rec,
+            gen_ema_classical,
+            [opt_dec, opt_disc],
+        ]

musika.py

@@ -0,0 +1,17 @@
+from parse_test import parse_args
+from models import Models_functions
+from utils import Utils_functions
+import os
+
+if __name__ == "__main__":
+
+    # parse args
+    args = parse_args()
+
+    # initialize networks
+    M = Models_functions(args)
+    models_ls_techno, models_ls_classical = M.get_networks()
+
+    # test musika
+    U = Utils_functions(args)
+    U.render_gradio(models_ls_techno, models_ls_classical, train=False)

parse_test.py

@@ -0,0 +1,165 @@
+import argparse
+from typing import Any
+import tensorflow as tf
+
+
+class EasyDict(dict):
+    def __getattr__(self, name: str) -> Any:
+        try:
+            return self[name]
+        except KeyError:
+            raise AttributeError(name)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        self[name] = value
+
+    def __delattr__(self, name: str) -> None:
+        del self[name]
+
+
+def params_args(args):
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--hop",
+        type=int,
+        default=256,
+        help="Hop size (window size = 4*hop)",
+    )
+    parser.add_argument(
+        "--mel_bins",
+        type=int,
+        default=256,
+        help="Mel bins in mel-spectrograms",
+    )
+    parser.add_argument(
+        "--sr",
+        type=int,
+        default=22050,
+        help="Sampling Rate",
+    )
+    parser.add_argument(
+        "--latlen",
+        type=int,
+        default=256,
+        help="Length of generated latent vectors",
+    )
+    parser.add_argument(
+        "--latdepth",
+        type=int,
+        default=64,
+        help="Depth of generated latent vectors",
+    )
+    parser.add_argument(
+        "--shape",
+        type=int,
+        default=128,
+        help="Length of spectrograms time axis",
+    )
+    parser.add_argument(
+        "--window",
+        type=int,
+        default=64,
+        help="Generator spectrogram window (must divide shape)",
+    )
+    parser.add_argument(
+        "--mu_rescale",
+        type=int,
+        default=-25.0,
+        help="Spectrogram mu used to normalize",
+    )
+    parser.add_argument(
+        "--sigma_rescale",
+        type=int,
+        default=75.0,
+        help="Spectrogram sigma used to normalize",
+    )
+    parser.add_argument(
+        "--load_path_techno",
+        type=str,
+        default="checkpoints/techno/",
+        help="Path of pretrained networks weights (techno)",
+    )
+    parser.add_argument(
+        "--load_path_classical",
+        type=str,
+        default="checkpoints/classical/",
+        help="Path of pretrained networks weights (classical)",
+    )
+    parser.add_argument(
+        "--dec_path_techno",
+        type=str,
+        default="checkpoints/techno/",
+        help="Path of pretrained decoders weights (techno)",
+    )
+    parser.add_argument(
+        "--dec_path_classical",
+        type=str,
+        default="checkpoints/classical/",
+        help="Path of pretrained decoders weights (classical)",
+    )
+    parser.add_argument(
+        "--testing",
+        type=bool,
+        default=True,
+        help="True if optimizers weight do not need to be loaded",
+    )
+    parser.add_argument(
+        "--cpu",
+        type=bool,
+        default=False,
+        help="True if you wish to use cpu",
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=bool,
+        default=True,
+        help="True if your GPU supports mixed precision",
+    )
+
+    tmp_args = parser.parse_args()
+
+    args.hop = tmp_args.hop
+    args.mel_bins = tmp_args.mel_bins
+    args.sr = tmp_args.sr
+    args.latlen = tmp_args.latlen
+    args.latdepth = tmp_args.latdepth
+    args.shape = tmp_args.shape
+    args.window = tmp_args.window
+    args.mu_rescale = tmp_args.mu_rescale
+    args.sigma_rescale = tmp_args.sigma_rescale
+    args.load_path_techno = tmp_args.load_path_techno
+    args.load_path_classical = tmp_args.load_path_classical
+    args.dec_path_techno = tmp_args.dec_path_techno
+    args.dec_path_classical = tmp_args.dec_path_classical
+    args.testing = tmp_args.testing
+    args.cpu = tmp_args.cpu
+    args.mixed_precision = tmp_args.mixed_precision
+
+    print()
+
+    args.datatype = tf.float32
+    gpuls = tf.config.list_physical_devices("GPU")
+    if len(gpuls) == 0 or args.cpu:
+        args.cpu = True
+        args.mixed_precision = False
+        tf.config.set_visible_devices([], "GPU")
+        print()
+        print("Using CPU...")
+        print()
+    if args.mixed_precision:
+        args.datatype = tf.float16
+        print()
+        print("Using GPU with mixed precision enabled...")
+        print()
+    if not args.mixed_precision and not args.cpu:
+        print()
+        print("Using GPU without mixed precision...")
+        print()
+
+    return args
+
+
+def parse_args():
+    args = EasyDict()
+    return params_args(args)

requirements.txt

@@ -0,0 +1,20 @@
+# This file may be used to create an environment using:
+# $ conda create --name <env> --file <this file>
+# platform: linux-64
+audioread==2.1.9
+gradio==2.5.2
+ipython==7.29.0
+librosa==0.8.1
+matplotlib==3.4.3
+numpy==1.20.3
+pillow==8.4.0
+protobuf==3.20.1rc1
+scikit-learn==1.0.1
+scipy==1.7.1
+seaborn==0.11.2
+soundfile==0.10.3.post1
+tensorboard==2.7.0
+tensorflow==2.7.0
+tensorflow-addons==0.15.0
+tensorflow-io==0.22.0
+tqdm==4.62.3

utils.py

@@ -0,0 +1,560 @@
+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)) at Johannes Kepler Universität Linz. Supervised by Jan Schlüter."
+
+        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", "Classical"],
+                    type="index",
+                    default="Techno",
+                    label="Music Genre to Generate (Brace yourself for very high levels of weirdness!)",
+                ),
+                gr.inputs.Radio(
+                    choices=["23 s", "115 s", "230 s"],
+                    type="index",
+                    default="115 s",
+                    label="Generated Music Length",
+                ),
+                gr.inputs.Slider(
+                    minimum=0,
+                    maximum=100,
+                    step=1,
+                    default=50,
+                    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",
+            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()
+        # iface.launch(share=True, enable_queue=True)
+        print("--------------------------------")
+        print("--------------------------------")
+        print("--------------------------------")
+        print("--------------------------------")
+        print("--------------------------------")