diff --git a/README.md b/README.md index 8b6c73d..adcd6cb 100644 --- a/README.md +++ b/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: ). 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` diff --git a/audio_separator/separator/audio_io.py b/audio_separator/separator/audio_io.py index d7d062d..efdde48 100644 --- a/audio_separator/separator/audio_io.py +++ b/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.""" diff --git a/audio_separator/separator/common_separator.py b/audio_separator/separator/common_separator.py index be20426..e515f4b 100644 --- a/audio_separator/separator/common_separator.py +++ b/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,8 +360,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 +392,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 +451,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: diff --git a/audio_separator/separator/separator.py b/audio_separator/separator/separator.py index ece2f10..e9d6236 100644 --- a/audio_separator/separator/separator.py +++ b/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, diff --git a/audio_separator/utils/cli.py b/audio_separator/utils/cli.py index 927b7ad..262986e 100755 --- a/audio_separator/utils/cli.py +++ b/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: ). 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, diff --git a/tests/unit/test_audio_output_contract.py b/tests/unit/test_audio_output_contract.py index cc63397..d6669ee 100644 --- a/tests/unit/test_audio_output_contract.py +++ b/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,6 +23,30 @@ 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.mark.parametrize(("output_format", "expected"), [("wav", "PCM_24"), ("flac", "PCM_24")]) +def test_auto_vorbis_falls_back_to_container_compatible_pcm(output_format, expected): + assert resolve_output_subtype("AUTO", "VORBIS", 24, output_format) == expected + + @pytest.fixture def common_separator(tmp_path): config = { diff --git a/tests/unit/test_bit_depth_writing.py b/tests/unit/test_bit_depth_writing.py index 1665f7d..f09e617 100644 --- a/tests/unit/test_bit_depth_writing.py +++ b/tests/unit/test_bit_depth_writing.py @@ -282,7 +282,36 @@ 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): + 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"]) - diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py index ebd419f..288de8d 100644 --- a/tests/unit/test_cli.py +++ b/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"] diff --git a/tests/unit/test_separator_api_compatibility.py b/tests/unit/test_separator_api_compatibility.py index 2875957..8b31732 100644 --- a/tests/unit/test_separator_api_compatibility.py +++ b/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"