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

@@ -7,9 +7,48 @@
 import tempfile
 
 import numpy as np
+import soundfile as sf
 
 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 sf.check_format(file_format.upper(), subtype):
+        return subtype
+
+    fallback = {16: "PCM_16", 24: "PCM_24", 32: "PCM_32"}.get(input_bit_depth, "PCM_16")
+    if file_format == "flac" and fallback == "PCM_32":
+        fallback = "PCM_24"
+    return fallback
+
 
 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
@@ -59,6 +64,7 @@ class CommonSeparator:
 
     def __init__(self, config):
 
+        """Initialize shared model, device, and output settings, validating the requested lossless subtype."""
         self.logger: Logger = config.get("logger")
         self.log_level: int = config.get("log_level")
 
@@ -78,6 +84,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,8 +361,18 @@ 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}")
 
-        # 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
+        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
+
+        output_bit_depth = {"PCM_16": 16, "PCM_24": 24, "PCM_32": 32}.get(output_subtype, 16)
         self.logger.info(f"Writing output with {output_bit_depth}-bit depth")
 
         # For pydub, we always convert to int16 for the AudioSegment creation
@@ -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, resolve_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,8 +23,36 @@
 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):
+    """Accept supported WAV subtypes case-insensitively and normalize their spelling."""
+    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):
+    """Reject unknown subtypes and incompatible container/subtype pairs before exporting audio."""
+    with pytest.raises(ValueError, match=message):
+        normalize_output_subtype(subtype, output_format)
+
+
+@pytest.mark.parametrize(("output_format", "expected"), [("wav", "PCM_24"), ("flac", "PCM_24")])
+def test_auto_vorbis_falls_back_to_container_compatible_pcm(output_format, expected):
+    """Choose compatible PCM output when AUTO encounters a VORBIS input subtype."""
+    assert resolve_output_subtype("AUTO", "VORBIS", 24, output_format) == expected
+
+
 @pytest.fixture
 def common_separator(tmp_path):
+    """Build a shared separator with mocked devices and a temporary WAV output directory."""
     config = {
         "logger": Mock(),
         "log_level": 20,

tests/unit/test_bit_depth_writing.py

@@ -282,7 +282,37 @@ 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)
+
+
+@requires_ffmpeg
+def test_explicit_16bit_output_overrides_24bit_input_metadata(temp_dir, mock_separator_config):
+    """Honor an explicit PCM_16 export even when the input metadata reports PCM_24."""
+    mock_separator_config["output_subtype"] = "PCM_16"
+    separator = CommonSeparator(mock_separator_config)
+    separator.input_bit_depth = 24
+    separator.input_subtype = "PCM_24"
+
+    samples = np.array([[0.25, -0.25], [0.5, -0.5]], dtype=np.float32)
+    separator.write_audio_pydub("explicit-16bit.wav", samples)
+
+    output_path = os.path.join(temp_dir, "explicit-16bit.wav")
+    assert sf.info(output_path).subtype == "PCM_16"
+
+
 if __name__ == "__main__":
     # Run tests with pytest
     pytest.main([__file__, "-v", "-s"])
-