musika clone

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Marco Pasini
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,13 +1,13 @@
 ---
-title: Musika Api
-emoji: 👀
-colorFrom: pink
-colorTo: red
+title: Musika
+emoji: 🎵
+colorFrom: purple
+colorTo: blue
 sdk: gradio
-sdk_version: 3.11.0
+sdk_version: 3.3.1
 app_file: app.py
 pinned: false
-license: mit
+license: cc-by-4.0
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

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_1, models_ls_2, models_ls_3 = M.get_networks()
+
+# test musika
+U = Utils_functions(args)
+U.render_gradio(models_ls_1, models_ls_2, models_ls_3, train=False)

checkpoints/ae/dec.h5

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

checkpoints/ae/dec2.h5

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

checkpoints/ae/enc.h5

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

checkpoints/ae/enc2.h5

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

checkpoints/metal/gen_ema.h5

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

checkpoints/misc/gen_ema.h5

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

checkpoints/techno/gen_ema.h5

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

layers.py

@@ -0,0 +1,163 @@
+import tensorflow as tf
+import tensorflow.python.keras.backend as K
+from tensorflow.python.eager import context
+from tensorflow.python.ops import (
+    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, datatype=tf.float32, **kwargs):
+        super(ConvSN2D, self).__init__(filters, kernel_size, **kwargs)
+        self.power_iterations = power_iterations
+        self.datatype = datatype
+
+    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(tf.keras.layers.Dense):
+    def __init__(self, datatype=tf.float32, **kwargs):
+        super(DenseSN, self).__init__(**kwargs)
+        self.datatype = datatype
+
+    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.datatype,
+        )
+
+    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 __init__(self, datatype=tf.float32, **kwargs):
+        super(AddNoise, self).__init__(**kwargs)
+        self.datatype = datatype
+
+    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.datatype,
+        )
+        output = inputs + self.b * rand
+        return output
+
+
+class PosEnc(tf.keras.layers.Layer):
+    def __init__(self, datatype=tf.float32, **kwargs):
+        super(PosEnc, self).__init__(**kwargs)
+        self.datatype = datatype
+
+    def call(self, inputs):
+        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.datatype)
+        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)

models.py

@@ -0,0 +1,783 @@
+import numpy as np
+import tensorflow as tf
+from tensorflow.python.keras.utils.layer_utils import count_params
+
+from layers import AddNoise
+
+
+class Models_functions:
+    def __init__(self, args):
+
+        self.args = args
+
+        if self.args.mixed_precision:
+            self.mixed_precision = tf.keras.mixed_precision
+            self.policy = tf.keras.mixed_precision.Policy("mixed_float16")
+            tf.keras.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", bnorm=True
+    ):
+
+        x = inp
+
+        bias = True
+        if bnorm:
+            bias = False
+
+        if upsample:
+            x = tf.keras.layers.Conv2DTranspose(
+                filters,
+                kernel_size=kernel_size,
+                strides=strides,
+                activation="linear",
+                padding=padding,
+                kernel_initializer=self.init,
+                use_bias=bias,
+            )(x)
+        else:
+            x = tf.keras.layers.Conv2D(
+                filters,
+                kernel_size=kernel_size,
+                strides=strides,
+                activation="linear",
+                padding=padding,
+                kernel_initializer=self.init,
+                use_bias=bias,
+            )(x)
+
+        if noise:
+            x = AddNoise(self.args.datatype)(x)
+
+        if bnorm:
+            x = tf.keras.layers.BatchNormalization()(x)
+
+        x = tf.keras.activations.swish(x)
+
+        return x
+
+    def pixel_shuffle(self, x, factor=2):
+        bs_dim, h_dim, w_dim, c_dim = tf.shape(x)[0], x.shape[1], x.shape[2], x.shape[3]
+        x = tf.reshape(x, [bs_dim, h_dim, w_dim, c_dim // factor, factor])
+        x = tf.transpose(x, [0, 1, 2, 4, 3])
+        return tf.reshape(x, [bs_dim, h_dim, w_dim * factor, c_dim // factor])
+
+    def adain(self, x, emb, name):
+        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,
+            name=name,
+        )(emb)
+        x = x / (tf.math.reduce_std(x, -2, keepdims=True) + 1e-5)
+        return x * emb
+
+    def conv_util_gen(
+        self,
+        inp,
+        filters,
+        kernel_size=(1, 9),
+        strides=(1, 1),
+        noise=False,
+        upsample=False,
+        emb=None,
+        se1=None,
+        name="0",
+    ):
+
+        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,
+                name=name + "c",
+            )(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,
+                name=name + "c",
+            )(x)
+
+        if noise:
+            x = AddNoise(self.args.datatype, name=name + "r")(x)
+
+        if emb is not None:
+            x = self.adain(x, emb, name=name + "ai")
+        else:
+            x = tf.keras.layers.BatchNormalization(name=name + "bn")(x)
+
+        x = tf.keras.activations.swish(x)
+
+        return x
+
+    def res_block_disc(self, inp, filters, kernel_size=(1, 3), kernel_size_2=None, strides=(1, 1), name="0"):
+
+        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,
+            name=name + "c0",
+        )(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,
+            name=name + "c1",
+        )(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,
+                name=name + "c3",
+            )(inp)
+
+        return x + inp
+
+    def build_encoder2(self):
+
+        inpf = tf.keras.layers.Input((1, self.args.shape, self.args.hop // 4))
+
+        inpfls = tf.split(inpf, 8, -2)
+        inpb = tf.concat(inpfls, 0)
+
+        g0 = self.conv_util(inpb, self.args.hop, kernel_size=(1, 3), strides=(1, 1), padding="same", bnorm=False)
+        g1 = self.conv_util(
+            g0, self.args.hop + self.args.hop // 2, kernel_size=(1, 3), strides=(1, 2), padding="valid", bnorm=False
+        )
+        g2 = self.conv_util(
+            g1, self.args.hop + self.args.hop // 2, kernel_size=(1, 3), strides=(1, 1), padding="same", bnorm=False
+        )
+        g3 = self.conv_util(g2, self.args.hop * 2, kernel_size=(1, 3), strides=(1, 2), padding="valid", bnorm=False)
+        g4 = self.conv_util(g3, self.args.hop * 2, kernel_size=(1, 3), strides=(1, 1), padding="same", bnorm=False)
+        g5 = self.conv_util(g4, self.args.hop * 3, kernel_size=(1, 3), strides=(1, 1), padding="valid", bnorm=False)
+        g5 = self.conv_util(g5, self.args.hop * 3, kernel_size=(1, 1), strides=(1, 1), padding="valid", bnorm=False)
+
+        g = tf.keras.layers.Conv2D(
+            self.args.latdepth,
+            kernel_size=(1, 1),
+            strides=1,
+            padding="valid",
+            kernel_initializer=self.init,
+            name="cbottle",
+            activation="tanh",
+        )(g5)
+
+        gls = tf.split(g, 8, 0)
+        g = tf.concat(gls, -2)
+        gls = tf.split(g, 2, -2)
+        g = tf.concat(gls, 0)
+
+        gf = tf.cast(g, tf.float32)
+
+        return tf.keras.Model(inpf, gf, name="ENC2")
+
+    def build_decoder2(self):
+
+        inpf = tf.keras.layers.Input((1, self.args.shape // 32, self.args.latdepth))
+
+        g = inpf
+
+        g = self.conv_util(
+            g, self.args.hop * 3, kernel_size=(1, 3), strides=(1, 1), upsample=False, noise=True, bnorm=False
+        )
+        g = self.conv_util(
+            g,
+            self.args.hop * 2 + self.args.hop // 2,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            upsample=True,
+            noise=True,
+            bnorm=False,
+        )
+        g = self.conv_util(
+            g,
+            self.args.hop * 2 + self.args.hop // 2,
+            kernel_size=(1, 3),
+            strides=(1, 1),
+            upsample=False,
+            noise=True,
+            bnorm=False,
+        )
+        g = self.conv_util(
+            g, self.args.hop * 2, kernel_size=(1, 4), strides=(1, 2), upsample=True, noise=True, bnorm=False
+        )
+        g = self.conv_util(
+            g, self.args.hop * 2, kernel_size=(1, 3), strides=(1, 1), upsample=False, noise=True, bnorm=False
+        )
+        g = self.conv_util(
+            g,
+            self.args.hop + self.args.hop // 2,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            upsample=True,
+            noise=True,
+            bnorm=False,
+        )
+        g = self.conv_util(g, self.args.hop, kernel_size=(1, 4), strides=(1, 2), upsample=True, noise=True, bnorm=False)
+
+        gf = tf.keras.layers.Conv2D(
+            self.args.hop // 4, kernel_size=(1, 1), strides=1, padding="same", kernel_initializer=self.init, name="cout"
+        )(g)
+
+        gfls = tf.split(gf, 2, 0)
+        gf = tf.concat(gfls, -2)
+
+        gf = tf.cast(gf, tf.float32)
+
+        return tf.keras.Model(inpf, gf, name="DEC2")
+
+    def build_encoder(self):
+
+        dim = ((4 * self.args.hop) // 2) + 1
+
+        inpf = tf.keras.layers.Input((dim, self.args.shape, 1))
+
+        ginp = tf.transpose(inpf, [0, 3, 2, 1])
+
+        g0 = self.conv_util(ginp, self.args.hop * 4, kernel_size=(1, 1), strides=(1, 1), padding="valid", bnorm=False)
+        g1 = self.conv_util(g0, self.args.hop * 4, kernel_size=(1, 1), strides=(1, 1), padding="valid", bnorm=False)
+        g2 = self.conv_util(g1, self.args.hop * 4, kernel_size=(1, 1), strides=(1, 1), padding="valid", bnorm=False)
+        g4 = self.conv_util(g2, self.args.hop * 4, kernel_size=(1, 1), strides=(1, 1), padding="valid", bnorm=False)
+        g5 = self.conv_util(g4, self.args.hop * 4, kernel_size=(1, 1), strides=(1, 1), padding="valid", bnorm=False)
+
+        g = tf.keras.layers.Conv2D(
+            self.args.hop // 4, kernel_size=(1, 1), strides=1, padding="valid", kernel_initializer=self.init
+        )(g5)
+
+        g = tf.keras.activations.tanh(g)
+
+        gls = tf.split(g, 2, -2)
+        g = tf.concat(gls, 0)
+
+        gf = tf.cast(g, tf.float32)
+
+        return tf.keras.Model(inpf, gf, name="ENC")
+
+    def build_decoder(self):
+
+        dim = ((4 * self.args.hop) // 2) + 1
+
+        inpf = tf.keras.layers.Input((1, self.args.shape // 2, self.args.hop // 4))
+
+        g = inpf
+
+        g0 = self.conv_util(g, self.args.hop * 3, kernel_size=(1, 3), strides=(1, 1), noise=True, bnorm=False)
+        g1 = self.conv_util(g0, self.args.hop * 3, kernel_size=(1, 3), strides=(1, 2), noise=True, bnorm=False)
+        g2 = self.conv_util(g1, self.args.hop * 2, kernel_size=(1, 3), strides=(1, 2), noise=True, bnorm=False)
+        g3 = self.conv_util(g2, self.args.hop, kernel_size=(1, 3), strides=(1, 2), noise=True, bnorm=False)
+        g = self.conv_util(g3, self.args.hop, kernel_size=(1, 3), strides=(1, 2), noise=True, bnorm=False)
+
+        g33 = self.conv_util(
+            g, self.args.hop, kernel_size=(1, 4), strides=(1, 2), upsample=True, noise=True, bnorm=False
+        )
+        g22 = self.conv_util(
+            g3, self.args.hop * 2, kernel_size=(1, 4), strides=(1, 2), upsample=True, noise=True, bnorm=False
+        )
+        g11 = self.conv_util(
+            g22 + g2, self.args.hop * 3, kernel_size=(1, 4), strides=(1, 2), upsample=True, noise=True, bnorm=False
+        )
+        g00 = self.conv_util(
+            g11 + g1, self.args.hop * 3, kernel_size=(1, 4), strides=(1, 2), upsample=True, noise=True, bnorm=False
+        )
+
+        g = tf.keras.layers.Conv2D(
+            dim, kernel_size=(1, 1), strides=(1, 1), kernel_initializer=self.init, padding="same"
+        )(g00 + g0)
+        gf = tf.clip_by_value(g, -1.0, 1.0)
+
+        g = self.conv_util(
+            g22, self.args.hop * 3, kernel_size=(1, 4), strides=(1, 2), upsample=True, noise=True, bnorm=False
+        )
+        g = self.conv_util(
+            g + g11, self.args.hop * 3, kernel_size=(1, 4), strides=(1, 2), upsample=True, noise=True, bnorm=False
+        )
+        g = tf.keras.layers.Conv2D(
+            dim, kernel_size=(1, 1), strides=(1, 1), kernel_initializer=self.init, padding="same"
+        )(g + g00)
+        pf = tf.clip_by_value(g, -1.0, 1.0)
+
+        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 tf.keras.Model(inpf, [s, p], name="DEC")
+
+    def build_critic(self):
+
+        sinp = tf.keras.layers.Input(shape=(1, self.args.latlen, self.args.latdepth * 2))
+
+        sf = tf.keras.layers.Conv2D(
+            self.args.base_channels * 3,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            activation="linear",
+            padding="same",
+            kernel_initializer=self.init,
+            name="1c",
+        )(sinp)
+        sf = tf.keras.layers.LeakyReLU(0.2)(sf)
+
+        sf = self.res_block_disc(sf, self.args.base_channels * 4, kernel_size=(1, 4), strides=(1, 2), name="2")
+
+        sf = self.res_block_disc(sf, self.args.base_channels * 5, kernel_size=(1, 4), strides=(1, 2), name="3")
+
+        sf = self.res_block_disc(sf, self.args.base_channels * 6, kernel_size=(1, 4), strides=(1, 2), name="4")
+
+        sf = self.res_block_disc(sf, self.args.base_channels * 7, kernel_size=(1, 4), strides=(1, 2), name="5")
+
+        if not self.args.small:
+            sf = self.res_block_disc(
+                sf, self.args.base_channels * 7, kernel_size=(1, 4), strides=(1, 2), kernel_size_2=(1, 1), name="6"
+            )
+
+        sf = tf.keras.layers.Conv2D(
+            self.args.base_channels * 7,
+            kernel_size=(1, 3),
+            strides=(1, 1),
+            activation="linear",
+            padding="same",
+            kernel_initializer=self.init,
+            name="7c",
+        )(sf)
+        sf = tf.keras.layers.LeakyReLU(0.2)(sf)
+
+        gf = tf.keras.layers.Dense(1, activation="linear", use_bias=True, kernel_initializer=self.init, name="7d")(
+            tf.keras.layers.Flatten()(sf)
+        )
+
+        gf = tf.cast(gf, tf.float32)
+
+        return tf.keras.Model(sinp, gf, name="C")
+
+    def build_generator(self):
+
+        dim = self.args.latdepth * 2
+
+        inpf = tf.keras.layers.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)
+        if not self.args.small:
+            inp6 = tf.keras.layers.AveragePooling2D((1, 2), padding="valid")(inp5)
+
+        if not self.args.small:
+            g = tf.keras.layers.Dense(
+                4 * (self.args.base_channels * 7),
+                activation="linear",
+                use_bias=True,
+                kernel_initializer=self.init,
+                name="00d",
+            )(tf.keras.layers.Flatten()(inp6))
+            g = tf.keras.layers.Reshape((1, 4, self.args.base_channels * 7))(g)
+            g = AddNoise(self.args.datatype, name="00n")(g)
+            g = self.adain(g, inp5, name="00ai")
+            g = tf.keras.activations.swish(g)
+        else:
+            g = tf.keras.layers.Dense(
+                4 * (self.args.base_channels * 7),
+                activation="linear",
+                use_bias=True,
+                kernel_initializer=self.init,
+                name="00d",
+            )(tf.keras.layers.Flatten()(inp5))
+            g = tf.keras.layers.Reshape((1, 4, self.args.base_channels * 7))(g)
+            g = AddNoise(self.args.datatype, name="00n")(g)
+            g = self.adain(g, inp4, name="00ai")
+            g = tf.keras.activations.swish(g)
+
+        if not self.args.small:
+            g1 = self.conv_util_gen(
+                g,
+                self.args.base_channels * 6,
+                kernel_size=(1, 4),
+                strides=(1, 2),
+                upsample=True,
+                noise=True,
+                emb=inp4,
+                name="0",
+            )
+            g1 = tf.math.sqrt(tf.cast(0.5, self.args.datatype)) * g1
+            g1 = self.conv_util_gen(
+                g1,
+                self.args.base_channels * 6,
+                kernel_size=(1, 4),
+                strides=(1, 1),
+                upsample=False,
+                noise=True,
+                emb=inp4,
+                name="1",
+            )
+            g1 = tf.math.sqrt(tf.cast(0.5, self.args.datatype)) * g1
+            g1 = g1 + tf.keras.layers.Conv2D(
+                g1.shape[-1],
+                kernel_size=(1, 1),
+                strides=1,
+                activation="linear",
+                padding="same",
+                kernel_initializer=self.init,
+                use_bias=True,
+                name="res1c",
+            )(self.pixel_shuffle(g))
+        else:
+            g1 = self.conv_util_gen(
+                g,
+                self.args.base_channels * 6,
+                kernel_size=(1, 1),
+                strides=(1, 1),
+                upsample=False,
+                noise=True,
+                emb=inp4,
+                name="0_small",
+            )
+            g1 = tf.math.sqrt(tf.cast(0.5, self.args.datatype)) * g1
+            g1 = self.conv_util_gen(
+                g1,
+                self.args.base_channels * 6,
+                kernel_size=(1, 1),
+                strides=(1, 1),
+                upsample=False,
+                noise=True,
+                emb=inp4,
+                name="1_small",
+            )
+            g1 = tf.math.sqrt(tf.cast(0.5, self.args.datatype)) * g1
+            g1 = g1 + tf.keras.layers.Conv2D(
+                g1.shape[-1],
+                kernel_size=(1, 1),
+                strides=1,
+                activation="linear",
+                padding="same",
+                kernel_initializer=self.init,
+                use_bias=True,
+                name="res1c_small",
+            )(g)
+
+        g2 = self.conv_util_gen(
+            g1,
+            self.args.base_channels * 5,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            upsample=True,
+            noise=True,
+            emb=inp3,
+            name="2",
+        )
+        g2 = tf.math.sqrt(tf.cast(0.5, self.args.datatype)) * g2
+        g2 = self.conv_util_gen(
+            g2,
+            self.args.base_channels * 5,
+            kernel_size=(1, 4),
+            strides=(1, 1),
+            upsample=False,
+            noise=True,
+            emb=inp3,
+            name="3",
+        )
+        g2 = tf.math.sqrt(tf.cast(0.5, self.args.datatype)) * g2
+        g2 = g2 + tf.keras.layers.Conv2D(
+            g2.shape[-1],
+            kernel_size=(1, 1),
+            strides=1,
+            activation="linear",
+            padding="same",
+            kernel_initializer=self.init,
+            use_bias=True,
+            name="res2c",
+        )(self.pixel_shuffle(g1))
+
+        g3 = self.conv_util_gen(
+            g2,
+            self.args.base_channels * 4,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            upsample=True,
+            noise=True,
+            emb=inp2,
+            name="4",
+        )
+        g3 = tf.math.sqrt(tf.cast(0.5, self.args.datatype)) * g3
+        g3 = self.conv_util_gen(
+            g3,
+            self.args.base_channels * 4,
+            kernel_size=(1, 4),
+            strides=(1, 1),
+            upsample=False,
+            noise=True,
+            emb=inp2,
+            name="5",
+        )
+        g3 = tf.math.sqrt(tf.cast(0.5, self.args.datatype)) * g3
+        g3 = g3 + tf.keras.layers.Conv2D(
+            g3.shape[-1],
+            kernel_size=(1, 1),
+            strides=1,
+            activation="linear",
+            padding="same",
+            kernel_initializer=self.init,
+            use_bias=True,
+            name="res3c",
+        )(self.pixel_shuffle(g2))
+
+        g4 = self.conv_util_gen(
+            g3,
+            self.args.base_channels * 3,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            upsample=True,
+            noise=True,
+            emb=inp1,
+            name="6",
+        )
+        g4 = tf.math.sqrt(tf.cast(0.5, self.args.datatype)) * g4
+        g4 = self.conv_util_gen(
+            g4,
+            self.args.base_channels * 3,
+            kernel_size=(1, 4),
+            strides=(1, 1),
+            upsample=False,
+            noise=True,
+            emb=inp1,
+            name="7",
+        )
+        g4 = tf.math.sqrt(tf.cast(0.5, self.args.datatype)) * g4
+        g4 = g4 + tf.keras.layers.Conv2D(
+            g4.shape[-1],
+            kernel_size=(1, 1),
+            strides=1,
+            activation="linear",
+            padding="same",
+            kernel_initializer=self.init,
+            use_bias=True,
+            name="res4c",
+        )(self.pixel_shuffle(g3))
+
+        g5 = self.conv_util_gen(
+            g4,
+            self.args.base_channels * 2,
+            kernel_size=(1, 4),
+            strides=(1, 2),
+            upsample=True,
+            noise=True,
+            emb=tf.expand_dims(tf.cast(inpb, dtype=self.args.datatype), -3),
+            name="8",
+        )
+
+        gf = tf.keras.layers.Conv2D(
+            dim, kernel_size=(1, 4), strides=(1, 1), kernel_initializer=self.init, padding="same", name="9c"
+        )(g5)
+
+        gfls = tf.split(gf, 2, 0)
+        gf = tf.concat(gfls, -2)
+
+        gf = tf.cast(gf, tf.float32)
+
+        return tf.keras.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()
+        gen_ema = self.build_generator()
+
+        switch = tf.Variable(-1.0, dtype=tf.float32)
+
+        if self.args.mixed_precision:
+            opt_disc = self.mixed_precision.LossScaleOptimizer(tf.keras.optimizers.Adam(0.0001, 0.5))
+            opt_dec = self.mixed_precision.LossScaleOptimizer(tf.keras.optimizers.Adam(0.0001, 0.5))
+        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")
+            enc.load_weights(self.args.dec_path + "/enc.h5")
+            enc2.load_weights(self.args.dec_path + "/enc2.h5")
+
+        else:
+            grad_vars = critic.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))
+                critic.load_weights(path + "/critic.h5")
+                gen.load_weights(path + "/gen.h5")
+                switch = tf.Variable(float(np.load(path + "/switch.npy", allow_pickle=True)), dtype=tf.float32)
+
+            gen_ema.load_weights(path + "/gen_ema.h5")
+            dec.load_weights(self.args.dec_path + "/dec.h5")
+            dec2.load_weights(self.args.dec_path + "/dec2.h5")
+            enc.load_weights(self.args.dec_path + "/enc.h5")
+            enc2.load_weights(self.args.dec_path + "/enc2.h5")
+
+        return (
+            critic,
+            gen,
+            enc,
+            dec,
+            enc2,
+            dec2,
+            gen_ema,
+            [opt_dec, opt_disc],
+            switch,
+        )
+
+    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()
+        gen_ema = self.build_generator()
+
+        switch = tf.Variable(-1.0, dtype=tf.float32)
+
+        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.5))
+            opt_dec = self.mixed_precision.LossScaleOptimizer(tf.keras.optimizers.Adam(0.0001, 0.5))
+        else:
+            opt_disc = tf.keras.optimizers.Adam(0.0001, 0.5)
+            opt_dec = tf.keras.optimizers.Adam(0.0001, 0.5)
+
+        return (
+            critic,
+            gen,
+            enc,
+            dec,
+            enc2,
+            dec2,
+            gen_ema,
+            [opt_dec, opt_disc],
+            switch,
+        )
+
+    def get_networks(self):
+        (
+            critic,
+            gen,
+            enc,
+            dec,
+            enc2,
+            dec2,
+            gen_ema_1,
+            [opt_dec, opt_disc],
+            switch,
+        ) = self.load(self.args.load_path_1, load_dec=False)
+        print(f"Networks loaded from {self.args.load_path_1}")
+
+        (
+            critic,
+            gen,
+            enc,
+            dec,
+            enc2,
+            dec2,
+            gen_ema_2,
+            [opt_dec, opt_disc],
+            switch,
+        ) = self.load(self.args.load_path_2, load_dec=False)
+        print(f"Networks loaded from {self.args.load_path_2}")
+
+        (
+            critic,
+            gen,
+            enc,
+            dec,
+            enc2,
+            dec2,
+            gen_ema_3,
+            [opt_dec, opt_disc],
+            switch,
+        ) = self.load(self.args.load_path_3, load_dec=False)
+        print(f"Networks loaded from {self.args.load_path_3}")
+
+        return (
+            (critic, gen, enc, dec, enc2, dec2, gen_ema_1, [opt_dec, opt_disc], switch),
+            (critic, gen, enc, dec, enc2, dec2, gen_ema_2, [opt_dec, opt_disc], switch),
+            (critic, gen, enc, dec, enc2, dec2, gen_ema_3, [opt_dec, opt_disc], switch),
+        )
+
+    def initialize_networks(self):
+
+        (
+            (critic, gen, enc, dec, enc2, dec2, gen_ema_1, [opt_dec, opt_disc], switch),
+            (critic, gen, enc, dec, enc2, dec2, gen_ema_2, [opt_dec, opt_disc], switch),
+            (critic, gen, enc, dec, enc2, dec2, gen_ema_3, [opt_dec, opt_disc], switch),
+        ) = self.get_networks()
+
+        print(f"Critic params: {count_params(critic.trainable_variables)}")
+        print(f"Generator params: {count_params(gen.trainable_variables)}")
+
+        return (
+            (critic, gen, enc, dec, enc2, dec2, gen_ema_1, [opt_dec, opt_disc], switch),
+            (critic, gen, enc, dec, enc2, dec2, gen_ema_2, [opt_dec, opt_disc], switch),
+            (critic, gen, enc, dec, enc2, dec2, gen_ema_3, [opt_dec, opt_disc], switch),
+        )

musika_test.py

@@ -0,0 +1,20 @@
+import os
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+from parse_test import parse_args
+from models import Models_functions
+from utils import Utils_functions
+
+if __name__ == "__main__":
+
+    # parse args
+    args = parse_args()
+
+    # initialize networks
+    M = Models_functions(args)
+    models_ls_1, models_ls_2, models_ls_3 = M.get_networks()
+
+    # test musika
+    U = Utils_functions(args)
+    U.render_gradio(models_ls_1, models_ls_2, models_ls_3, train=False)

parse_test.py

@@ -0,0 +1,196 @@
+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 str2bool(v):
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ("yes", "true", "t", "y", "1"):
+        return True
+    elif v.lower() in ("no", "false", "f", "n", "0"):
+        return False
+    else:
+        raise argparse.ArgumentTypeError("Boolean value expected.")
+
+
+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=44100,
+        help="Sampling Rate",
+    )
+    parser.add_argument(
+        "--small",
+        type=str2bool,
+        default=False,
+        help="If True, use model with shorter available context, useful for small datasets",
+    )
+    parser.add_argument(
+        "--latdepth",
+        type=int,
+        default=64,
+        help="Depth of generated latent vectors",
+    )
+    parser.add_argument(
+        "--coorddepth",
+        type=int,
+        default=64,
+        help="Dimension of latent coordinate and style random vectors",
+    )
+    parser.add_argument(
+        "--base_channels",
+        type=int,
+        default=128,
+        help="Base channels for generator and discriminator architectures",
+    )
+    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=float,
+        default=-25.0,
+        help="Spectrogram mu used to normalize",
+    )
+    parser.add_argument(
+        "--sigma_rescale",
+        type=float,
+        default=75.0,
+        help="Spectrogram sigma used to normalize",
+    )
+    parser.add_argument(
+        "--load_path_1",
+        type=str,
+        default="checkpoints/techno/",
+        help="Path of pretrained networks weights 1",
+    )
+    parser.add_argument(
+        "--load_path_2",
+        type=str,
+        default="checkpoints/metal/",
+        help="Path of pretrained networks weights 2",
+    )
+    parser.add_argument(
+        "--load_path_3",
+        type=str,
+        default="checkpoints/misc/",
+        help="Path of pretrained networks weights 3",
+    )
+    parser.add_argument(
+        "--dec_path",
+        type=str,
+        default="checkpoints/ae/",
+        help="Path of pretrained decoders weights",
+    )
+    parser.add_argument(
+        "--testing",
+        type=str2bool,
+        default=True,
+        help="True if optimizers weight do not need to be loaded",
+    )
+    parser.add_argument(
+        "--cpu",
+        type=str2bool,
+        default=False,
+        help="True if you wish to use cpu",
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str2bool,
+        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.small = tmp_args.small
+    args.latdepth = tmp_args.latdepth
+    args.coorddepth = tmp_args.coorddepth
+    args.base_channels = tmp_args.base_channels
+    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_1 = tmp_args.load_path_1
+    args.load_path_2 = tmp_args.load_path_2
+    args.load_path_3 = tmp_args.load_path_3
+    args.dec_path = tmp_args.dec_path
+    args.testing = tmp_args.testing
+    args.cpu = tmp_args.cpu
+    args.mixed_precision = tmp_args.mixed_precision
+
+    if args.small:
+        args.latlen = 128
+    else:
+        args.latlen = 256
+    args.coordlen = (args.latlen // 2) * 3
+
+    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,12 @@
+# This file may be used to create an environment using:
+# $ conda create --name <env> --file <this file>
+# platform: linux-64
+gradio==3.3.1
+librosa==0.8.1
+matplotlib==3.4.3
+numpy==1.20.3
+scipy==1.7.1
+tensorboard==2.10.0
+tensorflow==2.10.0
+tqdm==4.62.3
+pydub==0.25.1

utils.py

@@ -0,0 +1,689 @@
+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,
+        )
+
+        return (
+            np.clip(spec, -1.0, 1.0),
+            (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", "Death Metal (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.Image(label="Log-MelSpectrogram of Generated Audio (first 23 s)"),
+                gr.Audio(type="numpy", label="Generated Audio"),
+            ],
+            title="musika!",
+            description="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("--------------------------------")