Added AUDIO_SEPARATOR_MODEL_DIR env var to override model dir. Relaxed torch dependency again, with fix for demucs on torch 2.6. Improved inode utilization of test script.

Author: beveradb

audio_separator/separator/separator.py

@@ -123,15 +123,24 @@ def __init__(
             package_version = self.get_package_distribution("audio-separator").version
             self.logger.info(f"Separator version {package_version} instantiating with output_dir: {output_dir}, output_format: {output_format}")
 
-        self.model_file_dir = model_file_dir
-
         if output_dir is None:
             output_dir = os.getcwd()
             if not info_only:
                 self.logger.info("Output directory not specified. Using current working directory.")
 
         self.output_dir = output_dir
 
+        # Check for environment variable to override model_file_dir
+        env_model_dir = os.environ.get("AUDIO_SEPARATOR_MODEL_DIR")
+        if env_model_dir:
+            self.model_file_dir = env_model_dir
+            self.logger.info(f"Using model directory from AUDIO_SEPARATOR_MODEL_DIR env var: {self.model_file_dir}")
+            if not os.path.exists(self.model_file_dir):
+                raise FileNotFoundError(f"The specified model directory does not exist: {self.model_file_dir}")
+        else:
+            self.logger.info(f"Using model directory from model_file_dir parameter: {model_file_dir}")
+            self.model_file_dir = model_file_dir
+
         # Create the model directory if it does not exist
         os.makedirs(self.model_file_dir, exist_ok=True)
         os.makedirs(self.output_dir, exist_ok=True)

audio_separator/separator/uvr_lib_v5/demucs/states.py

@@ -40,7 +40,7 @@ def load_model(path_or_package, strict=False):
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")
             path = path_or_package
-            package = torch.load(path, "cpu")
+            package = torch.load(path, "cpu", weights_only=False)
     else:
         raise ValueError(f"Invalid type for {path_or_package}.")
 

audio_separator/utils/cli.py

@@ -41,7 +41,7 @@ def main():
     output_format_help = "Output format for separated files, any common format (default: %(default)s). Example: --output_format=MP3"
     output_bitrate_help = "Output bitrate for separated files, any ffmpeg-compatible bitrate (default: %(default)s). Example: --output_bitrate=320k"
     output_dir_help = "Directory to write output files (default: <current dir>). Example: --output_dir=/app/separated"
-    model_file_dir_help = "Model files directory (default: %(default)s). Example: --model_file_dir=/app/models"
+    model_file_dir_help = "Model files directory (default: %(default)s or AUDIO_SEPARATOR_MODEL_DIR env var if set). Example: --model_file_dir=/app/models"
     download_model_only_help = "Download a single model file only, without performing separation."
 
     io_params = parser.add_argument_group("Separation I/O Params")