Upload cog_predict.py

cog_predict.py

@@ -0,0 +1,144 @@
+import os
+import cog
+import tempfile
+import zipfile
+from pathlib import Path
+import argparse
+import data.utils
+import model.utils as model_utils
+from test import predict_song
+from model.waveunet import Waveunet
+
+
+class waveunetPredictor(cog.Predictor):
+    def setup(self):
+        """Init wave u net model"""
+        parser = argparse.ArgumentParser()
+        parser.add_argument(
+            "--instruments",
+            type=str,
+            nargs="+",
+            default=["bass", "drums", "other", "vocals"],
+            help='List of instruments to separate (default: "bass drums other vocals")',
+        )
+        parser.add_argument(
+            "--cuda", action="store_true", help="Use CUDA (default: False)"
+        )
+        parser.add_argument(
+            "--features",
+            type=int,
+            default=32,
+            help="Number of feature channels per layer",
+        )
+        parser.add_argument(
+            "--load_model",
+            type=str,
+            default="checkpoints/waveunet/model",
+            help="Reload a previously trained model",
+        )
+        parser.add_argument("--batch_size", type=int, default=4, help="Batch size")
+        parser.add_argument(
+            "--levels", type=int, default=6, help="Number of DS/US blocks"
+        )
+        parser.add_argument(
+            "--depth", type=int, default=1, help="Number of convs per block"
+        )
+        parser.add_argument("--sr", type=int, default=44100, help="Sampling rate")
+        parser.add_argument(
+            "--channels", type=int, default=2, help="Number of input audio channels"
+        )
+        parser.add_argument(
+            "--kernel_size",
+            type=int,
+            default=5,
+            help="Filter width of kernels. Has to be an odd number",
+        )
+        parser.add_argument(
+            "--output_size", type=float, default=2.0, help="Output duration"
+        )
+        parser.add_argument(
+            "--strides", type=int, default=4, help="Strides in Waveunet"
+        )
+        parser.add_argument(
+            "--conv_type",
+            type=str,
+            default="gn",
+            help="Type of convolution (normal, BN-normalised, GN-normalised): normal/bn/gn",
+        )
+        parser.add_argument(
+            "--res",
+            type=str,
+            default="fixed",
+            help="Resampling strategy: fixed sinc-based lowpass filtering or learned conv layer: fixed/learned",
+        )
+        parser.add_argument(
+            "--separate",
+            type=int,
+            default=1,
+            help="Train separate model for each source (1) or only one (0)",
+        )
+        parser.add_argument(
+            "--feature_growth",
+            type=str,
+            default="double",
+            help="How the features in each layer should grow, either (add) the initial number of features each time, or multiply by 2 (double)",
+        )
+        """
+        parser.add_argument('--input', type=str, default=str(input),
+                            help="Path to input mixture to be separated")
+        parser.add_argument('--output', type=str, default=out_path, help="Output path (same folder as input path if not set)")
+        """
+        args = parser.parse_args([])
+        self.args = args
+
+        num_features = (
+            [args.features * i for i in range(1, args.levels + 1)]
+            if args.feature_growth == "add"
+            else [args.features * 2 ** i for i in range(0, args.levels)]
+        )
+        target_outputs = int(args.output_size * args.sr)
+        self.model = Waveunet(
+            args.channels,
+            num_features,
+            args.channels,
+            args.instruments,
+            kernel_size=args.kernel_size,
+            target_output_size=target_outputs,
+            depth=args.depth,
+            strides=args.strides,
+            conv_type=args.conv_type,
+            res=args.res,
+            separate=args.separate,
+        )
+
+        if args.cuda:
+            self.model = model_utils.DataParallel(model)
+            print("move model to gpu")
+            self.model.cuda()
+
+        print("Loading model from checkpoint " + str(args.load_model))
+        state = model_utils.load_model(self.model, None, args.load_model, args.cuda)
+        print("Step", state["step"])
+
+    @cog.input("input", type=Path, help="audio mixture path")
+    def predict(self, input):
+        """Separate tracks from input mixture audio"""
+
+        out_path = Path(tempfile.mkdtemp())
+        zip_path = Path(tempfile.mkdtemp()) / "output.zip"
+
+        preds = predict_song(self.args, input, self.model)
+
+        out_names = []
+        for inst in preds.keys():
+            temp_n = os.path.join(
+                str(out_path), os.path.basename(str(input)) + "_" + inst + ".wav"
+            )
+            data.utils.write_wav(temp_n, preds[inst], self.args.sr)
+            out_names.append(temp_n)
+
+        with zipfile.ZipFile(str(zip_path), "w") as zf:
+            for i in out_names:
+                zf.write(str(i))
+
+        return zip_path