generation additions

MusicModel/decode.py

@@ -0,0 +1,21 @@
+import os
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+from parse.parse_decode 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)
+    M.download_networks()
+    models_ls = M.get_networks()
+
+    # encode samples
+    U = Utils_functions(args)
+    U.decode_path(models_ls)

MusicModel/encode.py

@@ -0,0 +1,24 @@
+import os
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+from parse.parse_encode import parse_args
+from models import Models_functions
+from utils_encode import UtilsEncode_functions
+
+if __name__ == "__main__":
+
+    # parse args
+    args = parse_args()
+
+    # initialize networks
+    M = Models_functions(args)
+    M.download_networks()
+    models_ls = M.get_networks()
+
+    # encode samples
+    U = UtilsEncode_functions(args)
+    if args.whole:
+        U.compress_whole_files(models_ls)
+    else:
+        U.compress_files(models_ls)

MusicModel/parse/parse_decode.py

@@ -0,0 +1,210 @@
+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(
+        "--max_lat_len",
+        type=int,
+        default=512,
+        help="Length of latent sequences: a random on-the-fly crop will be used for training",
+    )
+    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(
+        "--files_path",
+        type=str,
+        default="audio_samples/",
+        help="Path of compressed latent samples to decode",
+    )
+    parser.add_argument(
+        "--save_path",
+        type=str,
+        default="decoded_samples/",
+        help="Path where decoded audio files will be saved",
+    )
+    parser.add_argument(
+        "--dec_path",
+        type=str,
+        default="checkpoints/ae",
+        help="Path of pretrained decoders weights",
+    )
+    parser.add_argument(
+        "--load_path",
+        type=str,
+        default="None",
+        help="If not None, load models weights from this path",
+    )
+    parser.add_argument(
+        "--base_path",
+        type=str,
+        default="checkpoints",
+        help="Path where pretrained models are downloaded",
+    )
+    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.max_lat_len = tmp_args.max_lat_len
+    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.save_path = tmp_args.save_path
+    args.files_path = tmp_args.files_path
+    args.dec_path = tmp_args.dec_path
+    args.load_path = tmp_args.load_path
+    args.base_path = tmp_args.base_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)

MusicModel/parse/parse_encode.py

@@ -0,0 +1,217 @@
+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(
+        "--whole",
+        type=str2bool,
+        default=False,
+        help="If True, encode a single audio file to a single compressed encoding file of variable length",
+    )
+    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(
+        "--max_lat_len",
+        type=int,
+        default=512,
+        help="Length of latent sequences: a random on-the-fly crop will be used for training",
+    )
+    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(
+        "--files_path",
+        type=str,
+        default="audio_samples/",
+        help="Path of samples to encode",
+    )
+    parser.add_argument(
+        "--save_path",
+        type=str,
+        default="encoded_samples/",
+        help="Path where compressed representations will be saved",
+    )
+    parser.add_argument(
+        "--dec_path",
+        type=str,
+        default="checkpoints/ae",
+        help="Path of pretrained decoders weights",
+    )
+    parser.add_argument(
+        "--load_path",
+        type=str,
+        default="None",
+        help="If not None, load models weights from this path",
+    )
+    parser.add_argument(
+        "--base_path",
+        type=str,
+        default="checkpoints",
+        help="Path where pretrained models are downloaded",
+    )
+    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.whole = tmp_args.whole
+    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.max_lat_len = tmp_args.max_lat_len
+    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.save_path = tmp_args.save_path
+    args.files_path = tmp_args.files_path
+    args.dec_path = tmp_args.dec_path
+    args.load_path = tmp_args.load_path
+    args.base_path = tmp_args.base_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)

MusicModel/parse/parse_generate.py

@@ -0,0 +1,217 @@
+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(
+        "--num_samples",
+        type=int,
+        default=1,
+        help="Number of desired generated samples",
+    )
+    parser.add_argument(
+        "--seconds",
+        type=int,
+        default=120,
+        help="Length in seconds of generated samples",
+    )
+    parser.add_argument(
+        "--save_path",
+        type=str,
+        default="generations",
+        help="Path where to save generated samples",
+    )
+    parser.add_argument(
+        "--truncation",
+        type=float,
+        default=2.0,
+        help="Standard deviation of random vectors (truncation trick)",
+    )
+    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",
+        type=str,
+        default="checkpoints/techno/",
+        help="Path of pretrained networks weights",
+    )
+    parser.add_argument(
+        "--dec_path",
+        type=str,
+        default="checkpoints/ae/",
+        help="Path of pretrained decoders weights",
+    )
+    parser.add_argument(
+        "--base_path",
+        type=str,
+        default="checkpoints",
+        help="Path where pretrained models are downloaded",
+    )
+    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.num_samples = tmp_args.num_samples
+    args.seconds = tmp_args.seconds
+    args.save_path = tmp_args.save_path
+    args.truncation = tmp_args.truncation
+    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 = tmp_args.load_path
+    args.base_path = tmp_args.base_path
+    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)

MusicModel/parse/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",
+        type=str,
+        default="checkpoints/techno/",
+        help="Path of pretrained networks weights",
+    )
+    parser.add_argument(
+        "--dec_path",
+        type=str,
+        default="checkpoints/ae/",
+        help="Path of pretrained decoders weights",
+    )
+    parser.add_argument(
+        "--base_path",
+        type=str,
+        default="checkpoints",
+        help="Path where pretrained models are downloaded",
+    )
+    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",
+    )
+    parser.add_argument(
+        "--share_gradio",
+        type=str2bool,
+        default=False,
+        help="True if you wish to create a public URL for the Gradio interface",
+    )
+
+    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 = tmp_args.load_path
+    args.base_path = tmp_args.base_path
+    args.dec_path = tmp_args.dec_path
+    args.testing = tmp_args.testing
+    args.cpu = tmp_args.cpu
+    args.mixed_precision = tmp_args.mixed_precision
+    args.share_gradio = tmp_args.share_gradio
+
+    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)

MusicModel/utils_encode.py

@@ -0,0 +1,214 @@
+import os
+
+import numpy as np
+import tensorflow as tf
+from pydub import AudioSegment
+from glob import glob
+from tqdm import tqdm
+
+from utils import Utils_functions
+
+
+class UtilsEncode_functions:
+    def __init__(self, args):
+
+        self.args = args
+        self.U = Utils_functions(args)
+        self.paths = sorted(glob(self.args.files_path + "/*"))
+
+    def audio_generator(self):
+        for p in self.paths:
+            try:
+                tp, ext = os.path.splitext(p)
+                bname = os.path.basename(tp)
+                wvo = AudioSegment.from_file(p, format=ext[1:])
+                wvo = wvo.set_frame_rate(self.args.sr)
+                wvls = wvo.split_to_mono()
+                wvls = [s.get_array_of_samples() for s in wvls]
+                wv = np.array(wvls).T.astype(np.float32)
+                wv /= np.iinfo(wvls[0].typecode).max
+                yield np.squeeze(wv), bname
+            except Exception as e:
+                print(e)
+                print("Exception ignored! Continuing...")
+                pass
+
+    # def create_dataset(self):
+    #     self.ds = (
+    #         tf.data.Dataset.from_generator(
+    #             self.audio_generator, output_signature=(tf.TensorSpec(shape=(None, 2), dtype=tf.float32))
+    #         )
+    #         .prefetch(tf.data.experimental.AUTOTUNE)
+    #         .apply(tf.data.experimental.ignore_errors())
+    #     )
+
+    def compress_files(self, models_ls=None):
+        critic, gen, enc, dec, enc2, dec2, gen_ema, [opt_dec, opt_disc], switch = models_ls
+        # self.create_dataset()
+        os.makedirs(self.args.save_path, exist_ok=True)
+        c = 0
+        time_compression_ratio = 16  # TODO: infer time compression ratio
+        shape2 = self.args.shape
+        pbar = tqdm(self.audio_generator(), position=0, leave=True, total=len(self.paths))
+
+        for (wv,bname) in pbar:
+
+            try:
+
+                if wv.shape[0] > self.args.hop * self.args.shape * 2 + 3 * self.args.hop:
+
+                    split_limit = (
+                        5 * 60 * self.args.sr
+                    )  # split very long waveforms (> 5 minutes) and process separately to avoid out of memory errors
+
+                    nsplits = (wv.shape[0] // split_limit) + 1
+                    wvsplits = []
+                    for ns in range(nsplits):
+                        if wv.shape[0] - (ns * split_limit) > self.args.hop * self.args.shape * 2 + 3 * self.args.hop:
+                            wvsplits.append(wv[ns * split_limit : (ns + 1) * split_limit, :])
+
+                    for wv in wvsplits:
+
+                        wv = tf.image.random_crop(
+                            wv,
+                            size=[
+                                (((wv.shape[0] - (3 * self.args.hop)) // (self.args.shape * self.args.hop)))
+                                * self.args.shape
+                                * self.args.hop
+                                + 3 * self.args.hop,
+                                2,
+                            ],
+                        )
+
+                        chls = []
+                        for channel in range(2):
+
+                            x = wv[:, channel]
+                            x = tf.expand_dims(tf.transpose(self.U.wv2spec(x, hop_size=self.args.hop), (1, 0)), -1)
+                            ds = []
+                            num = x.shape[1] // self.args.shape
+                            rn = 0
+                            for i in range(num):
+                                ds.append(
+                                    x[:, rn + (i * self.args.shape) : rn + (i * self.args.shape) + self.args.shape, :]
+                                )
+                            del x
+                            ds = tf.convert_to_tensor(ds, dtype=tf.float32)
+                            lat = self.U.distribute_enc(ds, enc)
+                            del ds
+                            lat = tf.split(lat, lat.shape[0], 0)
+                            lat = tf.concat(lat, -2)
+                            lat = tf.squeeze(lat)
+
+                            switch = False
+                            if lat.shape[0] > (self.args.max_lat_len * time_compression_ratio):
+                                switch = True
+                                ds2 = []
+                                num2 = lat.shape[-2] // shape2
+                                rn2 = 0
+                                for j in range(num2):
+                                    ds2.append(lat[rn2 + (j * shape2) : rn2 + (j * shape2) + shape2, :])
+                                ds2 = tf.convert_to_tensor(ds2, dtype=tf.float32)
+                                lat = self.U.distribute_enc(tf.expand_dims(ds2, -3), enc2)
+                                del ds2
+                                lat = tf.split(lat, lat.shape[0], 0)
+                                lat = tf.concat(lat, -2)
+                                lat = tf.squeeze(lat)
+                                chls.append(lat)
+
+                        if lat.shape[0] > self.args.max_lat_len and switch:
+
+                            lat = tf.concat(chls, -1)
+
+                            del chls
+
+                            latc = lat[: (lat.shape[0] // self.args.max_lat_len) * self.args.max_lat_len, :]
+                            latc = tf.split(latc, latc.shape[0] // self.args.max_lat_len, 0)
+                            for el in latc:
+                                np.save(self.args.save_path + f"/{bname}_{c}.npy", el)
+                                c += 1
+                                pbar.set_postfix({"Saved Files": c})
+                            np.save(self.args.save_path + f"/{bname}_{c}.npy", lat[-self.args.max_lat_len :, :])
+                            c += 1
+                            pbar.set_postfix({"Saved Files": c})
+
+                            del lat
+                            del latc
+
+            except Exception as e:
+                print(e)
+                print("Exception ignored! Continuing...")
+                pass
+
+
+    def compress_whole_files(self, models_ls=None):
+        critic, gen, enc, dec, enc2, dec2, gen_ema, [opt_dec, opt_disc], switch = models_ls
+        # self.create_dataset()
+        os.makedirs(self.args.save_path, exist_ok=True)
+        c = 0
+        time_compression_ratio = 16  # TODO: infer time compression ratio
+        shape2 = self.args.shape
+        pbar = tqdm(self.audio_generator(), position=0, leave=True, total=len(self.paths))
+
+        for (wv,bname) in pbar:
+
+            try:
+
+                # wv_len_orig = wv.shape[0]
+
+                if wv.shape[0] > self.args.hop * self.args.shape * 2 + 3 * self.args.hop:
+
+                    rem = (wv.shape[0] - (3 * self.args.hop)) % (self.args.shape * self.args.hop)
+
+                    if rem != 0:
+                        wv = tf.concat([wv, tf.zeros([rem,2], dtype=tf.float32)], 0)
+
+                    chls = []
+                    for channel in range(2):
+
+                        x = wv[:, channel]
+                        x = tf.expand_dims(tf.transpose(self.U.wv2spec(x, hop_size=self.args.hop), (1, 0)), -1)
+                        ds = []
+                        num = x.shape[1] // self.args.shape
+                        rn = 0
+                        for i in range(num):
+                            ds.append(
+                                x[:, rn + (i * self.args.shape) : rn + (i * self.args.shape) + self.args.shape, :]
+                            )
+                        del x
+                        ds = tf.convert_to_tensor(ds, dtype=tf.float32)
+                        lat = self.U.distribute_enc(ds, enc)
+                        del ds
+                        lat = tf.split(lat, lat.shape[0], 0)
+                        lat = tf.concat(lat, -2)
+                        lat = tf.squeeze(lat)
+
+
+
+                        ds2 = []
+                        num2 = lat.shape[-2] // shape2
+                        rn2 = 0
+                        for j in range(num2):
+                            ds2.append(lat[rn2 + (j * shape2) : rn2 + (j * shape2) + shape2, :])
+                        ds2 = tf.convert_to_tensor(ds2, dtype=tf.float32)
+                        lat = self.U.distribute_enc(tf.expand_dims(ds2, -3), enc2)
+                        del ds2
+                        lat = tf.split(lat, lat.shape[0], 0)
+                        lat = tf.concat(lat, -2)
+                        lat = tf.squeeze(lat)
+                        chls.append(lat)
+
+                    lat = tf.concat(chls, -1)
+
+                    del chls
+
+                    np.save(self.args.save_path + f"/{bname}.npy", lat)
+                    c += 1
+                    pbar.set_postfix({"Saved Files": c})
+
+                    del lat
+
+            except Exception as e:
+                print(e)
+                print("Exception ignored! Continuing...")
+                pass

requirements.txt

@@ -1,7 +1,15 @@
-tensorflow==2.9.1
 gdown==4.4.0
 streamlit==1.10.0
 streamlit-tags==1.2.7
 torch==1.12.1
 scikit-learn==0.22
-transformers==4.24.0
+transformers==4.24.0
+
+librosa==0.8.1
+matplotlib==3.4.3
+numpy==1.23.5
+scipy==1.7.1
+tensorboard==2.10.0
+tensorflow==2.10.0
+tqdm==4.62.3
+pydub==0.25.1