update to musika 44khz

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

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

app.py

@@ -8,8 +8,8 @@ args = parse_args()
 
 # initialize networks
 M = Models_functions(args)
-models_ls_techno, models_ls_classical = M.get_networks()
+models_ls_1, models_ls_2, models_ls_3 = M.get_networks()
 
 # test musika
 U = Utils_functions(args)
-U.render_gradio(models_ls_techno, models_ls_classical, train=False)
+U.render_gradio(models_ls_1, models_ls_2, models_ls_3, train=False)

checkpoints/{classical → ae}/dec.h5

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

checkpoints/{techno → ae}/dec2.h5

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

checkpoints/{classical/dec2.h5 → ae/enc.h5}

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

checkpoints/{techno/dec.h5 → ae/enc2.h5}

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

checkpoints/{classical → metal}/gen_ema.h5

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1508fa1fad225f7c29ed64252b531e4462c97b0f0c80dd6d212014916f6261eb
-size 56431944
+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

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

layers.py

@@ -1,12 +1,7 @@
 import tensorflow as tf
-import tensorflow.keras.backend as K
-from tensorflow.keras.layers import Conv2D, Conv2DTranspose, Dense
+import tensorflow.python.keras.backend as K
 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,
@@ -19,9 +14,10 @@ def l2normalize(v, eps=1e-12):
 
 
 class ConvSN2D(tf.keras.layers.Conv2D):
-    def __init__(self, filters, kernel_size, power_iterations=1, **kwargs):
+    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)
@@ -70,7 +66,11 @@ class ConvSN2D(tf.keras.layers.Conv2D):
         return outputs
 
 
-class DenseSN(Dense):
+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)
 
@@ -79,7 +79,7 @@ class DenseSN(Dense):
             shape=tuple([1, self.kernel.shape.as_list()[-1]]),
             initializer=tf.initializers.RandomNormal(0, 1),
             trainable=False,
-            dtype=self.dtype,
+            dtype=self.datatype,
         )
 
     def compute_spectral_norm(self, W, new_u, W_shape):
@@ -116,6 +116,10 @@ class DenseSN(Dense):
 
 
 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=[
@@ -131,26 +135,24 @@ class AddNoise(tf.keras.layers.Layer):
             [tf.shape(inputs)[0], inputs.shape[1], inputs.shape[2], 1],
             mean=0.0,
             stddev=1.0,
-            dtype=self.dtype,
+            dtype=self.datatype,
         )
         output = inputs + self.b * rand
         return output
 
 
 class PosEnc(tf.keras.layers.Layer):
-    def __init__(self, **kwargs):
+    def __init__(self, datatype=tf.float32, **kwargs):
         super(PosEnc, self).__init__(**kwargs)
+        self.datatype = datatype
 
     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
-        )
+        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]
 
 
@@ -159,6 +161,3 @@ def flatten_hw(x, 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

@@ -1,60 +1,31 @@
 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
+from layers import 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.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",
-        bn=True,
+        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,
@@ -63,6 +34,7 @@ class Models_functions:
                 activation="linear",
                 padding=padding,
                 kernel_initializer=self.init,
+                use_bias=bias,
             )(x)
         else:
             x = tf.keras.layers.Conv2D(
@@ -72,19 +44,26 @@ class Models_functions:
                 activation="linear",
                 padding=padding,
                 kernel_initializer=self.init,
+                use_bias=bias,
             )(x)
 
         if noise:
-            x = AddNoise()(x)
+            x = AddNoise(self.args.datatype)(x)
 
-        if bn:
+        if bnorm:
             x = tf.keras.layers.BatchNormalization()(x)
 
         x = tf.keras.activations.swish(x)
 
         return x
 
-    def adain(self, x, emb):
+    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),
@@ -93,32 +72,11 @@ class Models_functions:
             padding="same",
             kernel_initializer=self.init,
             use_bias=True,
+            name=name,
         )(emb)
-        x = x / (tf.math.reduce_std(x, -2, keepdims=True) + 1e-7)
+        x = x / (tf.math.reduce_std(x, -2, keepdims=True) + 1e-5)
         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,
@@ -129,6 +87,7 @@ class Models_functions:
         upsample=False,
         emb=None,
         se1=None,
+        name="0",
     ):
 
         x = inp
@@ -142,6 +101,7 @@ class Models_functions:
                 padding="same",
                 kernel_initializer=self.init,
                 use_bias=True,
+                name=name + "c",
             )(x)
         else:
             x = tf.keras.layers.Conv2D(
@@ -152,24 +112,22 @@ class Models_functions:
                 padding="same",
                 kernel_initializer=self.init,
                 use_bias=True,
+                name=name + "c",
             )(x)
 
         if noise:
-            x = AddNoise()(x)
+            x = AddNoise(self.args.datatype, name=name + "r")(x)
 
         if emb is not None:
-            x = self.adain(x, emb)
+            x = self.adain(x, emb, name=name + "ai")
         else:
-            x = tf.keras.layers.BatchNormalization()(x)
-
-        x1 = tf.keras.activations.swish(x)
+            x = tf.keras.layers.BatchNormalization(name=name + "bn")(x)
 
-        if se1 is not None:
-            x1 = x1 * se1
+        x = tf.keras.activations.swish(x)
 
-        return x1
+        return x
 
-    def res_block_disc(self, inp, filters, kernel_size=(1, 3), kernel_size_2=None, strides=(1, 1)):
+    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
@@ -181,6 +139,7 @@ class Models_functions:
             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
@@ -191,6 +150,7 @@ class Models_functions:
             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
@@ -207,28 +167,32 @@ class Models_functions:
                 padding="same",
                 kernel_initializer=self.init,
                 use_bias=False,
+                name=name + "c3",
             )(inp)
 
         return x + inp
 
     def build_encoder2(self):
 
-        dim = 128
-
-        inpf = Input((1, self.args.shape, dim))
+        inpf = tf.keras.layers.Input((1, self.args.shape, self.args.hop // 4))
 
-        inpfls = tf.split(inpf, 16, -2)
+        inpfls = tf.split(inpf, 8, -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")
+        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(
-            64,
+            self.args.latdepth,
             kernel_size=(1, 1),
             strides=1,
             padding="valid",
@@ -237,50 +201,61 @@ class Models_functions:
             activation="tanh",
         )(g5)
 
-        gls = tf.split(g, 16, 0)
+        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 Model(inpf, gf, name="ENC2")
 
-    def build_decoder2(self):
+        return tf.keras.Model(inpf, gf, name="ENC2")
 
-        dim = 128
-        bottledim = 64
+    def build_decoder2(self):
 
-        inpf = Input((1, self.args.shape // 16, bottledim))
+        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,
-            512 + 128 + 128,
+            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,
-            512 + 128 + 128,
+            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, 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)
+        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(
-            dim,
-            kernel_size=(1, 1),
-            strides=1,
-            padding="same",
-            activation="tanh",
-            kernel_initializer=self.init,
+            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)
@@ -288,131 +263,77 @@ class Models_functions:
 
         gf = tf.cast(gf, tf.float32)
 
-        return Model(inpf, gf, name="DEC2")
+        return tf.keras.Model(inpf, gf, name="DEC2")
 
     def build_encoder(self):
 
         dim = ((4 * self.args.hop) // 2) + 1
 
-        inpf = Input((dim, self.args.shape, 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 * 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",
-        )
+        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(
-            128,
-            kernel_size=(1, 1),
-            strides=1,
-            padding="valid",
-            kernel_initializer=self.init,
-        )(g)
-        gb = tf.keras.activations.tanh(g)
+            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)
 
-        gbls = tf.split(gb, 2, -2)
-        gb = tf.concat(gbls, 0)
+        gf = tf.cast(g, tf.float32)
 
-        gb = tf.cast(gb, tf.float32)
-        return Model(inpf, gb, name="ENC")
+        return tf.keras.Model(inpf, gf, name="ENC")
 
     def build_decoder(self):
 
         dim = ((4 * self.args.hop) // 2) + 1
 
-        inpf = Input((1, self.args.shape // 2, 128))
+        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, 1), strides=(1, 1), noise=True)
+        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)
 
-        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,
+        g33 = self.conv_util(
+            g, self.args.hop, kernel_size=(1, 4), strides=(1, 2), upsample=True, noise=True, bnorm=False
         )
-        g = self.conv_util(
-            g2,
-            self.args.hop + self.args.hop // 4,
-            kernel_size=(1, 3),
-            strides=(1, 2),
-            noise=True,
+        g22 = self.conv_util(
+            g3, 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 + self.args.hop // 2,
-            kernel_size=(1, 4),
-            strides=(1, 2),
-            upsample=True,
-            noise=True,
+        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(
-            g + g2,
-            self.args.hop * 2,
-            kernel_size=(1, 4),
-            strides=(1, 2),
-            upsample=True,
-            noise=True,
+            g22, self.args.hop * 3, kernel_size=(1, 4), strides=(1, 2), upsample=True, noise=True, bnorm=False
         )
         g = self.conv_util(
-            g + g1,
-            self.args.hop * 3,
-            kernel_size=(1, 4),
-            strides=(1, 2),
-            upsample=True,
-            noise=True,
+            g + g11, self.args.hop * 3, kernel_size=(1, 4), strides=(1, 2), upsample=True, noise=True, bnorm=False
         )
-
-        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)
+        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)
@@ -426,128 +347,60 @@ class Models_functions:
         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")
+        return tf.keras.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)
+        sinp = tf.keras.layers.Input(shape=(1, self.args.latlen, self.args.latdepth * 2))
 
         sf = tf.keras.layers.Conv2D(
-            256 + 128,
-            kernel_size=(1, 3),
+            self.args.base_channels * 3,
+            kernel_size=(1, 4),
             strides=(1, 2),
             activation="linear",
             padding="same",
             kernel_initializer=self.init,
-        )(sf)
+            name="1c",
+        )(sinp)
         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))
+        sf = self.res_block_disc(sf, self.args.base_channels * 4, kernel_size=(1, 4), strides=(1, 2), name="2")
 
-        gf = tf.cast(gf, tf.float32)
-        sfo = tf.cast(sfo, tf.float32)
+        sf = self.res_block_disc(sf, self.args.base_channels * 5, kernel_size=(1, 4), strides=(1, 2), name="3")
 
-        return Model(sinp, [gf, sfo], name="C")
+        sf = self.res_block_disc(sf, self.args.base_channels * 6, kernel_size=(1, 4), strides=(1, 2), name="4")
 
-    def build_critic_rec(self):
+        sf = self.res_block_disc(sf, self.args.base_channels * 7, kernel_size=(1, 4), strides=(1, 2), name="5")
 
-        sinp = Input(shape=(1, self.args.latlen // 64, 512 + 512))
+        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"
+            )
 
-        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),
+        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.Conv2D(
-            dim,
-            kernel_size=(1, 1),
-            strides=(1, 1),
-            activation="tanh",
-            padding="same",
-            kernel_initializer=self.init,
-        )(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 Model(sinp, gf, name="CR")
+        return tf.keras.Model(sinp, gf, name="C")
 
     def build_generator(self):
 
         dim = self.args.latdepth * 2
 
-        inpf = Input((self.args.latlen, 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)
@@ -558,112 +411,213 @@ class Models_functions:
         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)
+        if not self.args.small:
+            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)
+        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)
 
-        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,
+            self.args.base_channels * 5,
             kernel_size=(1, 4),
             strides=(1, 2),
             upsample=True,
             noise=True,
             emb=inp3,
+            name="2",
         )
-        g2b = self.conv_util_gen(
+        g2 = tf.math.sqrt(tf.cast(0.5, self.args.datatype)) * g2
+        g2 = self.conv_util_gen(
             g2,
-            512 + 128,
-            kernel_size=(1, 3),
+            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(
-            g2b,
-            256 + 256,
+            g2,
+            self.args.base_channels * 4,
             kernel_size=(1, 4),
             strides=(1, 2),
             upsample=True,
             noise=True,
             emb=inp2,
-            se1=self.se_layer(g, 256 + 256),
+            name="4",
         )
+        g3 = tf.math.sqrt(tf.cast(0.5, self.args.datatype)) * g3
         g3 = self.conv_util_gen(
             g3,
-            256 + 256,
-            kernel_size=(1, 3),
+            self.args.base_channels * 4,
+            kernel_size=(1, 4),
             strides=(1, 1),
             upsample=False,
             noise=True,
             emb=inp2,
-            se1=self.se_layer(g1, 256 + 256),
+            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,
-            256 + 128,
+            self.args.base_channels * 3,
             kernel_size=(1, 4),
             strides=(1, 2),
             upsample=True,
             noise=True,
             emb=inp1,
-            se1=self.se_layer(g2, 256 + 128),
+            name="6",
         )
+        g4 = tf.math.sqrt(tf.cast(0.5, self.args.datatype)) * g4
         g4 = self.conv_util_gen(
             g4,
-            256 + 128,
-            kernel_size=(1, 3),
+            self.args.base_channels * 3,
+            kernel_size=(1, 4),
             strides=(1, 1),
             upsample=False,
             noise=True,
             emb=inp1,
-            se1=self.se_layer(g2b, 256 + 128),
+            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,
-            256,
+            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, 1),
-            strides=(1, 1),
-            kernel_initializer=self.init,
-            padding="same",
-            activation="tanh",
+            dim, kernel_size=(1, 4), strides=(1, 1), kernel_initializer=self.init, padding="same", name="9c"
         )(g5)
 
         gfls = tf.split(gf, 2, 0)
@@ -671,7 +625,7 @@ class Models_functions:
 
         gf = tf.cast(gf, tf.float32)
 
-        return Model(inpf, gf, name="GEN")
+        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):
@@ -681,12 +635,13 @@ class Models_functions:
         dec = self.build_decoder()
         enc2 = self.build_encoder2()
         dec2 = self.build_decoder2()
-        critic_rec = self.build_critic_rec()
         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.9))
-            opt_dec = self.mixed_precision.LossScaleOptimizer(tf.keras.optimizers.Adam(0.0001, 0.9))
+            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)
@@ -694,9 +649,11 @@ class Models_functions:
         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 + critic_rec.trainable_weights
+            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))
 
@@ -707,16 +664,15 @@ class Models_functions:
             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")
+                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(path + "/dec.h5")
-            dec2.load_weights(path + "/dec2.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,
@@ -725,9 +681,9 @@ class Models_functions:
             dec,
             enc2,
             dec2,
-            critic_rec,
             gen_ema,
             [opt_dec, opt_disc],
+            switch,
         )
 
     def build(self):
@@ -737,18 +693,19 @@ class Models_functions:
         dec = self.build_decoder()
         enc2 = self.build_encoder2()
         dec2 = self.build_decoder2()
-        critic_rec = self.build_critic_rec()
         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.9))
-            opt_dec = self.mixed_precision.LossScaleOptimizer(tf.keras.optimizers.Adam(0.0001, 0.9))
+            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)
+            opt_disc = tf.keras.optimizers.Adam(0.0001, 0.5)
+            opt_dec = tf.keras.optimizers.Adam(0.0001, 0.5)
 
         return (
             critic,
@@ -757,9 +714,9 @@ class Models_functions:
             dec,
             enc2,
             dec2,
-            critic_rec,
             gen_ema,
             [opt_dec, opt_disc],
+            switch,
         )
 
     def get_networks(self):
@@ -767,65 +724,60 @@ class Models_functions:
             critic,
             gen,
             enc,
-            dec_techno,
+            dec,
             enc2,
-            dec2_techno,
-            critic_rec,
-            gen_ema_techno,
+            dec2,
+            gen_ema_1,
             [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}")
+            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_classical,
+            dec,
             enc2,
-            dec2_classical,
-            critic_rec,
-            gen_ema_classical,
+            dec2,
+            gen_ema_2,
             [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}")
+            switch,
+        ) = self.load(self.args.load_path_2, load_dec=False)
+        print(f"Networks loaded from {self.args.load_path_2}")
 
-        return [critic, gen, enc, dec_techno, enc2, dec2_techno, critic_rec, gen_ema_techno, [opt_dec, opt_disc]], [
+        (
             critic,
             gen,
             enc,
-            dec_classical,
+            dec,
             enc2,
-            dec2_classical,
-            critic_rec,
-            gen_ema_classical,
+            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_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()
+        (
+            (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"Generator params: {count_params(gen_ema_techno.trainable_variables)}")
-        print(f"Decoder params: {count_params(dec_techno.trainable_variables+dec2_techno.trainable_variables)}")
+        print(f"Critic params: {count_params(critic.trainable_variables)}")
+        print(f"Generator params: {count_params(gen.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],
-        ]
+        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.py → musika_test.py

@@ -1,7 +1,10 @@
+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
-import os
 
 if __name__ == "__main__":
 
@@ -10,8 +13,8 @@ if __name__ == "__main__":
 
     # initialize networks
     M = Models_functions(args)
-    models_ls_techno, models_ls_classical = M.get_networks()
+    models_ls_1, models_ls_2, models_ls_3 = M.get_networks()
 
     # test musika
     U = Utils_functions(args)
-    U.render_gradio(models_ls_techno, models_ls_classical, train=False)
+    U.render_gradio(models_ls_1, models_ls_2, models_ls_3, train=False)

parse_test.py

@@ -17,6 +17,17 @@ class EasyDict(dict):
         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()
 
@@ -35,14 +46,14 @@ def params_args(args):
     parser.add_argument(
         "--sr",
         type=int,
-        default=22050,
+        default=44100,
         help="Sampling Rate",
     )
     parser.add_argument(
-        "--latlen",
-        type=int,
-        default=256,
-        help="Length of generated latent vectors",
+        "--small",
+        type=str2bool,
+        default=False,
+        help="If True, use model with shorter available context, useful for small datasets",
     )
     parser.add_argument(
         "--latdepth",
@@ -50,6 +61,18 @@ def params_args(args):
         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,
@@ -64,55 +87,55 @@ def params_args(args):
     )
     parser.add_argument(
         "--mu_rescale",
-        type=int,
+        type=float,
         default=-25.0,
         help="Spectrogram mu used to normalize",
     )
     parser.add_argument(
         "--sigma_rescale",
-        type=int,
+        type=float,
         default=75.0,
         help="Spectrogram sigma used to normalize",
     )
     parser.add_argument(
-        "--load_path_techno",
+        "--load_path_1",
         type=str,
         default="checkpoints/techno/",
-        help="Path of pretrained networks weights (techno)",
+        help="Path of pretrained networks weights 1",
     )
     parser.add_argument(
-        "--load_path_classical",
+        "--load_path_2",
         type=str,
-        default="checkpoints/classical/",
-        help="Path of pretrained networks weights (classical)",
+        default="checkpoints/metal/",
+        help="Path of pretrained networks weights 2",
     )
     parser.add_argument(
-        "--dec_path_techno",
+        "--load_path_3",
         type=str,
-        default="checkpoints/techno/",
-        help="Path of pretrained decoders weights (techno)",
+        default="checkpoints/misc/",
+        help="Path of pretrained networks weights 3",
     )
     parser.add_argument(
-        "--dec_path_classical",
+        "--dec_path",
         type=str,
-        default="checkpoints/classical/",
-        help="Path of pretrained decoders weights (classical)",
+        default="checkpoints/ae/",
+        help="Path of pretrained decoders weights",
     )
     parser.add_argument(
         "--testing",
-        type=bool,
+        type=str2bool,
         default=True,
         help="True if optimizers weight do not need to be loaded",
     )
     parser.add_argument(
         "--cpu",
-        type=bool,
+        type=str2bool,
         default=False,
         help="True if you wish to use cpu",
     )
     parser.add_argument(
         "--mixed_precision",
-        type=bool,
+        type=str2bool,
         default=True,
         help="True if your GPU supports mixed precision",
     )
@@ -122,20 +145,28 @@ def params_args(args):
     args.hop = tmp_args.hop
     args.mel_bins = tmp_args.mel_bins
     args.sr = tmp_args.sr
-    args.latlen = tmp_args.latlen
+    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_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.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

requirements.txt

@@ -1,20 +1,12 @@
 # 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
+gradio==3.3.1
 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
+tensorboard==2.10.0
+tensorflow==2.10.0
 tqdm==4.62.3
+pydub==0.25.1

utils.py

@@ -1,21 +1,13 @@
-import io
 import os
-import random
 import time
+import datetime
 from glob import glob
-
-import IPython
+from tqdm import tqdm
 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
 
@@ -27,14 +19,18 @@ class Utils_functions:
 
         melmat = tf.signal.linear_to_mel_weight_matrix(
             num_mel_bins=args.mel_bins,
-            num_spectrogram_bins=(4 * args.hop) // 2 + 1,
+            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,
+            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))
@@ -149,7 +145,17 @@ class Utils_functions:
         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):
+    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]
@@ -161,14 +167,13 @@ class Utils_functions:
                 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),
+            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):
+    def distribute_enc(self, x, model, bs=32):
         outls = []
         if isinstance(x, list):
             bdim = x[0].shape[0]
@@ -185,9 +190,9 @@ class Utils_functions:
                 res = tf.concat(resls, -2)
                 outls.append(res)
 
-        return np.concatenate(outls, 0)
+        return tf.concat(outls, 0)
 
-    def distribute_dec(self, x, model, bs=64):
+    def distribute_dec(self, x, model, bs=32):
         outls = []
         bdim = x.shape[0]
         for i in range(((bdim - 2) // bs) + 1):
@@ -196,12 +201,11 @@ class Utils_functions:
             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),
+        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):
+    def distribute_dec2(self, x, model, bs=32):
         outls = []
         bdim = x.shape[0]
         for i in range(((bdim - 2) // bs) + 1):
@@ -210,32 +214,105 @@ class Utils_functions:
             inp1 = tf.concat(inpls, 0)
             outls.append(model(inp1, training=False))
 
-        return np.concatenate(outls, 0)
+        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, 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)
+        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.latlen + 1, axis=-2)[:, :-1, :]
-        rr = tf.linspace(noisec, noiser, self.args.latlen + 1, axis=-2)
+        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)
-        return tf.image.random_crop(noisetot, [1, self.args.latlen, 64 + 64])
+        noisetot = self.center_coordinate(noisetot)
+        return self.crop_coordinate(noisetot)
 
     def generate_example_stereo(self, models_ls):
-        (critic, gen, enc, dec, enc2, dec2, critic_rec, gen_ema, [opt_dec, opt_disc],) = 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] // 16, -2)
+        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,
+                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)
@@ -273,14 +350,28 @@ class Utils_functions:
 
         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,),
+            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, 256), [1, 0],)), -2,
+                np.array(
+                    tf.transpose(
+                        self.wv2spec_hop((abwv2[:, 0] + abwv2[:, 1]) / 2.0, 80.0, self.args.hop * 2),
+                        [1, 0],
+                    )
+                ),
+                -2,
             ),
             cmap=None,
         )
@@ -288,7 +379,13 @@ class Utils_functions:
         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,
+                np.array(
+                    tf.transpose(
+                        self.wv2spec_hop((abwv3[:, 0] + abwv3[:, 1]) / 2.0, 80.0, self.args.hop * 2),
+                        [1, 0],
+                    )
+                ),
+                -2,
             ),
             cmap=None,
         )
@@ -296,7 +393,13 @@ class Utils_functions:
         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,
+                np.array(
+                    tf.transpose(
+                        self.wv2spec_hop((abwv4[:, 0] + abwv4[:, 1]) / 2.0, 80.0, self.args.hop * 2),
+                        [1, 0],
+                    )
+                ),
+                -2,
             ),
             cmap=None,
         )
@@ -304,18 +407,24 @@ class Utils_functions:
         axs[3].set_title("Generated4")
         # plt.show()
         plt.savefig(f"{path}/output.png")
+        plt.close()
 
-    # Save in training loop
     def save_end(
-        self, epoch, gloss, closs, mloss, models_ls=None, n_save=3, save_path="checkpoints",
+        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
+        (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!_-{str(gloss)[:9]}-{str(closs)[:9]}-{str(mloss)[:9]}"
+            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")
-            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")
@@ -324,84 +433,51 @@ class Utils_functions:
             # 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=64):
+    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 np.concatenate(outls, 0)
+        return tf.concat(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
+    def generate_waveform(self, inp, gen_ema, dec, dec2, batch_size=64):
 
-        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)
+        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)
 
-    def stfunc(self, genre, z, var, models_ls_techno, models_ls_classical):
+        chls = []
+        for channel in range(2):
 
-        (
-            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
+            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)
 
-        var = 0.01 + (3.5 * (var / 100.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)
 
-        if z == 0:
-            fac = 1
-        elif z == 1:
-            fac = 5
-        else:
-            fac = 10
+        return np.clip(np.squeeze(np.stack(chls, -1)), -1.0, 1.0)
 
-        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
+    def decode_waveform(self, lat, dec, dec2, batch_size=64):
 
-        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)
+        lat = lat[:, :, : (lat.shape[-2] // 8) * 8, :]
+        abls = tf.split(lat, lat.shape[-2] // 8, -2)
         abi = tf.concat(abls, 0)
 
         chls = []
@@ -410,25 +486,135 @@ class Utils_functions:
             ab = self.distribute_dec2(
                 abi[:, :, :, channel * self.args.latdepth : channel * self.args.latdepth + self.args.latdepth],
                 dec2,
-                bs=128,
+                bs=batch_size,
             )
-            # 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)
+            ab_m, ab_p = self.distribute_dec(ab, dec, bs=batch_size)
             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!"
+        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,
         )
 
-        abwvc = np.clip(np.squeeze(np.stack(chls, -1)), -1.0, 1.0)
+        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, 256),
+                    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],
                 )
             ),
@@ -436,46 +622,53 @@ class Utils_functions:
         )
 
         return (
-            spec,
+            np.clip(spec, -1.0, 1.0),
             (self.args.sr, np.int16(abwvc * 32767.0)),
         )
 
-    def render_gradio(self, models_ls_techno, models_ls_classical, train=True):
-        article_text = "Original work by Marco Pasini ([Twitter](https://twitter.com/marco_ppasini)) and Jan Schlüter at Johannes Kepler Universität Linz."
+    def 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(x, y, z):
-            return self.stfunc(x, y, z, models_ls_techno, models_ls_classical)
+        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.inputs.Radio(
-                    choices=["Techno/Experimental", "Classical"],
+                gr.Radio(
+                    choices=["Techno/Experimental", "Death Metal (finetuned)", "Misc"],
                     type="index",
-                    default="Classical",
+                    value="Techno/Experimental",
                     label="Music Genre to Generate",
                 ),
-                gr.inputs.Radio(
-                    choices=["23s", "1m 58s", "3m 57s"], type="index", default="1m 58s", label="Generated Music Length",
+                gr.Radio(
+                    choices=durations,
+                    type="index",
+                    value=durations_default,
+                    label="Generated Music Length",
                 ),
-                gr.inputs.Slider(
-                    minimum=0,
-                    maximum=100,
-                    step=1,
-                    default=25,
-                    label="Stability[left]/Variety[right] Tradeoff (Truncation Trick)",
+                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.outputs.Image(label="Log-MelSpectrogram of Generated Audio (first 23 sec)"),
-                gr.outputs.Audio(type="numpy", label="Generated Audio"),
+                gr.Image(label="Log-MelSpectrogram of Generated Audio (first 23 s)"),
+                gr.Audio(type="numpy", label="Generated Audio"),
             ],
-            allow_screenshot=False,
             title="musika!",
-            description="Blazingly Fast Stereo Waveform Music Generation of Arbitrary Length. Be patient and enjoy the weirdness!",
+            description="Blazingly Fast 44.1 kHz Stereo Waveform Music Generation of Arbitrary Length. Be patient and enjoy the weirdness!",
             article=article_text,
-            layout="vertical",
-            theme="huggingface",
         )
 
         print("--------------------------------")