Add support for custom model

base.sds

@@ -1,4 +1,4 @@
-%Srdr-1.0
+%Srdr-1.2
 
 '
 ' Surrounder base scene file
@@ -13,7 +13,7 @@ Chs: CHANNELS_5_1; 16; 44100
 Lay: LAYOUT_5_1
 
 ' Splitting into stems
-Sep: STEMS_ALL
+Sep: STEMS_ALL; MODEL_DEFAULT
 
 ' Original version of the song in the front channels
 Map: STEMS_PIANO_L; 0; 1

constants.py

@@ -11,9 +11,12 @@ def value_to_constant(value, filter_str=None):
 
     return keys[0]
 
+Version = "0.3.0"
+
 Versions = [
     "1.0",
-    "1.1"
+    "1.1",
+    "1.2"
 ]
 
 Channels = [
@@ -154,5 +157,15 @@ def value_to_constant(value, filter_str=None):
     'OUTPUT_AC3': 2,
     'OUTPUT_EAC3': 3,
     'OUTPUT_AC4': 4,
-    'OUTPUT_AAC': 5
+    'OUTPUT_AAC': 5,
+
+    # AI models
+    'MODEL_LEGACY': "htdemucs_6s",
+    'MODEL_STANDARD': "htdemucs",
+    'MODEL_STANDARD_LG': "htdemucs_ft",
+    'MODEL_STANDARD_XL': "htdemucs_mmi",
+    'MODEL_DEFAULT': "mdx_extra",
+    'MODEL_BETTER_MD': "mdx",
+    'MODEL_BETTER_SM': "mdx_extra_q",
+    'MODEL_BETTER_XS': "mdx_q",
 }

main.py

@@ -46,12 +46,16 @@
     exit(1)
 
 version = lines[0].split("-")[1]
-print(f"Surrounder Scene version: {version}")
 
 if version not in Versions:
     print(f"This version of Surrounder does not support this Scene file. Supported versions: {', '.join(Versions)}")
     exit(1)
 
+if Versions.index(version) == len(Versions) - 1:
+    print(f"Surrounder {Version}, using Scene specification version {version} (native)")
+else:
+    print(f"Surrounder {Version}, using Scene specification version {version} (compatibility)")
+
 operations = []
 
 for line in lines[1:]:
@@ -175,15 +179,26 @@
                     f"./srdr_work/channels/{i}.wav"
                 ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
         case "Sep":
-            if len(parameters) != 1:
-                print(f"Expected 1 parameter but got {len(parameters)}.")
-                exit(2)
+            if version == "1.0" or version == "1.1":
+                if len(parameters) != 1:
+                    print(f"Expected 1 parameter but got {len(parameters)}.")
+                    exit(2)
+            else:
+                if len(parameters) != 2:
+                    print(f"Expected 2 parameters but got {len(parameters)}.")
+                    exit(2)
 
             if not os.path.exists(f"./srdr_work/input.wav"):
                 print(f"Separation is not ready: not configured.")
                 exit(2)
 
-            args = ["--float32", "-n", "htdemucs_6s", "-o", "./srdr_work/stems_tmp"]
+            model = "htdemucs_6s"
+
+            if version != "1.0" and version != "1.1":
+                model = parameters[1]
+
+            print(f"Using model {model}")
+            args = ["--float32", "-n", model, "-o", "./srdr_work/stems_tmp"]
 
             try:
                 stem_format = int(parameters[0])
@@ -229,7 +244,7 @@
                 import demucs.separate
                 print("Separating using machine learning, this might take a while.")
                 demucs.separate.main(args)
-                os.rename("./srdr_work/stems_tmp/htdemucs_6s/input", "./srdr_work/stems")
+                os.rename(f"./srdr_work/stems_tmp/{model}/input", "./srdr_work/stems")
                 shutil.rmtree('./srdr_work/stems_tmp')
             elif stem_format == -1:
                 os.mkdir("./srdr_work/stems")