Add lossless output subtype selection

README.md

@@ -520,7 +520,7 @@ Presets are defined in `audio_separator/ensemble_presets.json` — contributions
 
 ```sh
 usage: audio-separator [-h] [-v] [-d] [-e] [-l] [--log_level LOG_LEVEL] [--list_filter LIST_FILTER] [--list_limit LIST_LIMIT] [--list_format {pretty,json}] [-m MODEL_FILENAME] [--output_format OUTPUT_FORMAT]
-                       [--output_bitrate OUTPUT_BITRATE] [--output_dir OUTPUT_DIR] [--model_file_dir MODEL_FILE_DIR] [--download_model_only] [--invert_spect] [--normalization NORMALIZATION]
+                       [--output_bitrate OUTPUT_BITRATE] [--output_subtype {AUTO,PCM_16,PCM_24,PCM_32,FLOAT}] [--output_dir OUTPUT_DIR] [--model_file_dir MODEL_FILE_DIR] [--download_model_only] [--invert_spect] [--normalization NORMALIZATION]
                        [--amplification AMPLIFICATION] [--single_stem SINGLE_STEM] [--sample_rate SAMPLE_RATE] [--use_soundfile] [--use_autocast | --use_native_fp16] [--use_torch_compile] [--use_directml] [--custom_output_names CUSTOM_OUTPUT_NAMES]
                        [--mdx_segment_size MDX_SEGMENT_SIZE] [--mdx_overlap MDX_OVERLAP] [--mdx_batch_size MDX_BATCH_SIZE] [--mdx_hop_length MDX_HOP_LENGTH] [--mdx_enable_denoise] [--vr_batch_size VR_BATCH_SIZE]
                        [--vr_window_size VR_WINDOW_SIZE] [--vr_aggression VR_AGGRESSION] [--vr_enable_tta] [--vr_high_end_process] [--vr_enable_post_process]
@@ -551,6 +551,7 @@ Separation I/O Params:
   -m MODEL_FILENAME, --model_filename MODEL_FILENAME     Model to use for separation (default: model_bs_roformer_ep_317_sdr_12.9755.yaml). Example: -m 2_HP-UVR.pth
   --output_format OUTPUT_FORMAT                          Output format for separated files, any common format (default: FLAC). Example: --output_format=MP3
   --output_bitrate OUTPUT_BITRATE                        Output bitrate for separated files, any ffmpeg-compatible bitrate (default: None). Example: --output_bitrate=320k
+  --output_subtype {AUTO,PCM_16,PCM_24,PCM_32,FLOAT}     Lossless WAV/FLAC subtype. AUTO follows the input where supported (default: AUTO).
   --output_dir OUTPUT_DIR                                Directory to write output files (default: <current dir>). Example: --output_dir=/app/separated
   --model_file_dir MODEL_FILE_DIR                        Model files directory (default: /tmp/audio-separator-models/). Example: --model_file_dir=/app/models
   --download_model_only                                  Download a single model file only, without performing separation.
@@ -744,6 +745,7 @@ You can also rename specific stems:
 - **`invert_using_spec`:** (Optional) Flag to invert using spectrogram. `Default: False`
 - **`sample_rate`:** (Optional) Set the sample rate of the output audio. `Default: 44100`
 - **`use_soundfile`:** (Optional) Use soundfile for output writing, can solve OOM issues, especially on longer audio.
+- **`output_subtype`:** (Optional) Choose `PCM_16`, `PCM_24`, `PCM_32`, or `FLOAT` for lossless output. `AUTO` preserves the input subtype where the WAV/FLAC container supports it. `PCM_32` and `FLOAT` require WAV.
 - **`use_autocast`:** (Optional) Use PyTorch autocast when the loaded model and device support it. Mutually exclusive with `use_native_fp16=True`. `Default: False`
 - **`use_native_fp16`:** (Optional) Convert a verified model to native float16 inference. Currently supported for MelBand RoFormer and BS-RoFormer on MPS and CUDA. Mutually exclusive with `use_autocast=True`; unsupported combinations warn and continue in float32. `Default: False`
 - **`use_torch_compile`:** (Optional) Compile verified repeated model blocks. This can be combined with float32 or autocast for MelBand RoFormer and BS-RoFormer on CPU, MPS, and CUDA, and with native float16 on MPS and CUDA. A fresh compiler cache can make the first run slower; unsupported combinations warn and continue in eager mode. `Default: False`

audio_separator/separator/audio_io.py

@@ -10,6 +10,40 @@
 
 from audio_separator.separator.exceptions import AudioExportError, InvalidAudioDataError
 
+OUTPUT_SUBTYPES = ("AUTO", "PCM_16", "PCM_24", "PCM_32", "FLOAT")
+
+
+def normalize_output_subtype(output_subtype, output_format):
+    """Validate and normalize an explicit lossless output subtype."""
+    subtype = str(output_subtype or "AUTO").upper()
+    if subtype not in OUTPUT_SUBTYPES:
+        choices = ", ".join(OUTPUT_SUBTYPES)
+        raise ValueError(f"output_subtype must be one of: {choices}")
+
+    file_format = str(output_format or "WAV").lower()
+    if subtype != "AUTO" and file_format not in ("wav", "flac"):
+        raise ValueError("output_subtype can only be set for WAV or FLAC output")
+    if file_format == "flac" and subtype in ("PCM_32", "FLOAT"):
+        raise ValueError(f"FLAC output does not support {subtype}; use PCM_24 or WAV")
+    return subtype
+
+
+def resolve_output_subtype(requested_subtype, input_subtype, input_bit_depth, output_format):
+    """Resolve AUTO to an input-compatible subtype supported by the container."""
+    subtype = normalize_output_subtype(requested_subtype, output_format)
+    if subtype != "AUTO":
+        return subtype
+
+    file_format = str(output_format or "WAV").lower()
+    if input_subtype:
+        subtype = str(input_subtype).upper()
+    else:
+        subtype = {16: "PCM_16", 24: "PCM_24", 32: "PCM_32"}.get(input_bit_depth, "PCM_16")
+
+    if file_format == "flac" and subtype in ("PCM_32", "FLOAT", "DOUBLE"):
+        return "PCM_24"
+    return subtype
+
 
 def validate_audio_source(stem_source):
     """Return audio as an array after validating mono/stereo frame layout."""

audio_separator/separator/common_separator.py

@@ -9,7 +9,12 @@
 import torch
 from pydub import AudioSegment
 import soundfile as sf
-from audio_separator.separator.audio_io import atomic_output_path, validate_audio_source
+from audio_separator.separator.audio_io import (
+    atomic_output_path,
+    normalize_output_subtype,
+    resolve_output_subtype,
+    validate_audio_source,
+)
 from audio_separator.separator.exceptions import AudioExportError, InvalidAudioDataError
 from audio_separator.separator.uvr_lib_v5 import spec_utils
 from audio_separator.separator.execution_policy import FP32, resolve_execution_policy
@@ -78,6 +83,7 @@ def __init__(self, config):
         self.output_dir = config.get("output_dir")
         self.output_format = config.get("output_format")
         self.output_bitrate = config.get("output_bitrate")
+        self.output_subtype = normalize_output_subtype(config.get("output_subtype", "AUTO"), self.output_format)
 
         # Functional options which are applicable to all architectures and the user may tweak to affect the output
         self.normalization_threshold = config.get("normalization_threshold")
@@ -354,6 +360,17 @@ def write_audio_pydub(self, stem_path: str, stem_source):
         self.logger.debug(f"Audio data shape before processing: {stem_source.shape}")
         self.logger.debug(f"Data type before conversion: {stem_source.dtype}")
 
+        file_format = stem_path.lower().split(".")[-1]
+        output_subtype = resolve_output_subtype(self.output_subtype, self.input_subtype, self.input_bit_depth, file_format)
+
+        # Pydub starts from int16 samples, so higher-depth lossless output must
+        # write the model's floating-point samples directly through soundfile.
+        if file_format in ("wav", "flac") and output_subtype != "PCM_16":
+            with atomic_output_path(stem_path, "soundfile") as temp_path:
+                sf.write(temp_path, stem_source, self.sample_rate, subtype=output_subtype)
+            self.logger.debug(f"Exported audio file successfully to {stem_path} with subtype {output_subtype}")
+            return
+
         # Determine bit depth for output (use input bit depth if available, otherwise default to 16)
         output_bit_depth = self.input_bit_depth if self.input_bit_depth is not None else 16
         self.logger.info(f"Writing output with {output_bit_depth}-bit depth")
@@ -376,9 +393,6 @@ def write_audio_pydub(self, stem_path: str, stem_source):
         except Exception as e:
             raise AudioExportError(f"Failed to create audio for {stem_path} with pydub: {e}", path=stem_path, backend="pydub") from e
 
-        # Determine file format based on the file extension
-        file_format = stem_path.lower().split(".")[-1]
-
         # For m4a files, specify mp4 as the container format as the extension doesn't match the format name
         if file_format == "m4a":
             file_format = "mp4"
@@ -438,26 +452,9 @@ def write_audio_soundfile(self, stem_path: str, stem_source):
             except Exception as e:
                 raise AudioExportError(f"Failed to prepare output directory for {stem_path}: {e}", path=stem_path, backend="soundfile") from e
 
-        # Determine the subtype based on the input audio's bit depth
-        output_subtype = None
-        if self.input_subtype:
-            output_subtype = self.input_subtype
-            self.logger.info(f"Using input subtype for output: {output_subtype}")
-        elif self.input_bit_depth:
-            # Map bit depth to subtype
-            if self.input_bit_depth == 16:
-                output_subtype = 'PCM_16'
-            elif self.input_bit_depth == 24:
-                output_subtype = 'PCM_24'
-            elif self.input_bit_depth == 32:
-                output_subtype = 'PCM_32'
-            else:
-                output_subtype = 'PCM_16'  # Default fallback
-            self.logger.info(f"Using output subtype based on bit depth: {output_subtype}")
-        else:
-            # Default to PCM_16 if no bit depth info available
-            output_subtype = 'PCM_16'
-            self.logger.warning("No bit depth info available, defaulting to PCM_16")
+        file_format = stem_path.lower().split(".")[-1]
+        output_subtype = resolve_output_subtype(self.output_subtype, self.input_subtype, self.input_bit_depth, file_format)
+        self.logger.info(f"Using output subtype: {output_subtype}")
 
         # Correctly interleave stereo channels if needed
         if stem_source.ndim == 2 and stem_source.shape[1] == 2:

audio_separator/separator/separator.py

@@ -23,7 +23,7 @@
 import torch.amp.autocast_mode as autocast_mode
 import onnxruntime as ort
 from tqdm import tqdm
-from audio_separator.separator.audio_io import atomic_output_path, validate_audio_source
+from audio_separator.separator.audio_io import atomic_output_path, normalize_output_subtype, validate_audio_source
 from audio_separator.separator.ensembler import Ensembler
 from audio_separator.separator.exceptions import AudioExportError, BatchSeparationError, InvalidAudioDataError
 from audio_separator.separator.execution_policy import AUTOCAST, FP32, NATIVE_FP16
@@ -88,6 +88,7 @@ class Separator:
         use_autocast (bool): Use PyTorch autocast when the loaded model and device support it.
         use_torch_compile (bool): Compile verified repeated model blocks when supported.
         use_native_fp16 (bool): Convert a verified model to native float16 inference when supported.
+        output_subtype (str): Lossless output subtype: AUTO, PCM_16, PCM_24, PCM_32, or FLOAT.
 
     MDX Architecture Specific Attributes:
         hop_length (int): The hop length for STFT.
@@ -146,6 +147,7 @@ def __init__(
         info_only=False,
         use_torch_compile=False,
         use_native_fp16=False,
+        output_subtype="AUTO",
     ):
         """Initialize the separator."""
         if use_autocast and use_native_fp16:
@@ -202,6 +204,7 @@ def __init__(
 
         if self.output_format is None:
             self.output_format = "WAV"
+        self.output_subtype = normalize_output_subtype(output_subtype, self.output_format)
 
         self.normalization_threshold = normalization_threshold
         if normalization_threshold <= 0 or normalization_threshold > 1:
@@ -961,6 +964,7 @@ def load_model(self, model_filename="model_bs_roformer_ep_317_sdr_12.9755.ckpt",
                 "model_data": model_data,
                 "output_format": self.output_format,
                 "output_bitrate": self.output_bitrate,
+                "output_subtype": self.output_subtype,
                 "output_dir": self.output_dir,
                 "normalization_threshold": self.normalization_threshold,
                 "amplification_threshold": self.amplification_threshold,

audio_separator/utils/cli.py

@@ -41,6 +41,7 @@ def main():
     extra_models_help = "Additional models for ensembling. Requires -m for the primary model. Example: --extra_models model2.onnx model3.ckpt"
     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_subtype_help = "Lossless WAV/FLAC subtype. AUTO follows the input where supported (default: %(default)s)."
     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 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."
@@ -50,6 +51,9 @@ def main():
     io_params.add_argument("--extra_models", nargs="+", default=None, help=extra_models_help)
     io_params.add_argument("--output_format", default="FLAC", help=output_format_help)
     io_params.add_argument("--output_bitrate", default=None, help=output_bitrate_help)
+    io_params.add_argument(
+        "--output_subtype", choices=["AUTO", "PCM_16", "PCM_24", "PCM_32", "FLOAT"], default="AUTO", help=output_subtype_help
+    )
     io_params.add_argument("--output_dir", default=None, help=output_dir_help)
     io_params.add_argument("--model_file_dir", default="/tmp/audio-separator-models/", help=model_file_dir_help)
     io_params.add_argument("--download_model_only", action="store_true", help=download_model_only_help)
@@ -260,6 +264,7 @@ def main():
         output_dir=args.output_dir,
         output_format=args.output_format,
         output_bitrate=args.output_bitrate,
+        output_subtype=args.output_subtype,
         normalization_threshold=args.normalization,
         amplification_threshold=args.amplification,
         output_single_stem=args.single_stem,

tests/unit/test_audio_output_contract.py

@@ -13,7 +13,7 @@
 import torch
 from pydub import AudioSegment
 
-from audio_separator.separator.audio_io import validate_audio_source
+from audio_separator.separator.audio_io import normalize_output_subtype, validate_audio_source
 from audio_separator.separator.architectures.demucs_separator import DemucsSeparator
 from audio_separator.separator.common_separator import CommonSeparator
 from audio_separator.separator.uvr_lib_v5 import spec_utils
@@ -23,6 +23,25 @@
 requires_ffmpeg = pytest.mark.skipif(shutil.which("ffmpeg") is None, reason="FFmpeg is required for pydub encoding")
 
 
+@pytest.mark.parametrize("subtype", ["AUTO", "PCM_16", "PCM_24", "PCM_32", "FLOAT"])
+def test_wav_accepts_supported_output_subtypes(subtype):
+    assert normalize_output_subtype(subtype.lower(), "WAV") == subtype
+
+
+@pytest.mark.parametrize(
+    ("subtype", "output_format", "message"),
+    [
+        ("PCM_24", "MP3", "WAV or FLAC"),
+        ("PCM_32", "FLAC", "does not support"),
+        ("FLOAT", "FLAC", "does not support"),
+        ("PCM_20", "WAV", "must be one of"),
+    ],
+)
+def test_invalid_output_subtype_combinations_fail_early(subtype, output_format, message):
+    with pytest.raises(ValueError, match=message):
+        normalize_output_subtype(subtype, output_format)
+
+
 @pytest.fixture
 def common_separator(tmp_path):
     config = {

tests/unit/test_bit_depth_writing.py

@@ -282,7 +282,22 @@ def test_write_16bit_with_soundfile(temp_dir, mock_separator_config):
     print("✅ Test passed: 16-bit audio written correctly with soundfile")
 
 
+def test_explicit_24bit_output_preserves_model_precision(temp_dir, mock_separator_config):
+    """An explicit high-depth output must bypass pydub's int16 conversion."""
+    mock_separator_config["output_subtype"] = "PCM_24"
+    separator = CommonSeparator(mock_separator_config)
+    separator.input_bit_depth = 16
+    separator.input_subtype = "PCM_16"
+
+    samples = np.array([[0.123456, -0.123456], [0.234567, -0.234567]], dtype=np.float32)
+    separator.write_audio_pydub("explicit-24bit.wav", samples)
+
+    output_path = os.path.join(temp_dir, "explicit-24bit.wav")
+    assert sf.info(output_path).subtype == "PCM_24"
+    decoded, _ = sf.read(output_path, dtype="float32")
+    np.testing.assert_allclose(decoded, samples, atol=1 / 2**23)
+
+
 if __name__ == "__main__":
     # Run tests with pytest
     pytest.main([__file__, "-v", "-s"])
-

tests/unit/test_cli.py

@@ -36,6 +36,7 @@ def common_expected_args():
         "output_dir": None,
         "output_format": "FLAC",
         "output_bitrate": None,
+        "output_subtype": "AUTO",
         "normalization_threshold": 0.9,
         "amplification_threshold": 0.0,
         "output_single_stem": None,
@@ -197,6 +198,18 @@ def test_cli_output_format_argument(common_expected_args):
             mock_separator.assert_called_once_with(**expected_args)
 
 
+def test_cli_output_subtype_argument(common_expected_args):
+    test_args = ["cli.py", "test_audio.wav", "--output_subtype=PCM_24"]
+    with patch("sys.argv", test_args):
+        with patch("audio_separator.separator.Separator") as mock_separator:
+            mock_separator.return_value.separate.return_value = ["output_file.wav"]
+            main()
+
+    expected_args = common_expected_args.copy()
+    expected_args["output_subtype"] = "PCM_24"
+    mock_separator.assert_called_once_with(**expected_args)
+
+
 # Test using normalization_threshold argument
 def test_cli_normalization_threshold_argument(common_expected_args):
     test_args = ["cli.py", "test_audio.mp3", "--normalization=0.75"]

tests/unit/test_separator_api_compatibility.py

@@ -33,6 +33,8 @@ def test_execution_options_are_appended_to_constructor_signature():
         "info_only",
         "use_torch_compile",
         "use_native_fp16",
+        "output_subtype",
     ]
     assert inspect.signature(Separator.__init__).parameters["use_torch_compile"].default is False
     assert inspect.signature(Separator.__init__).parameters["use_native_fp16"].default is False
+    assert inspect.signature(Separator.__init__).parameters["output_subtype"].default == "AUTO"