splifft 0.0.2__py3-none-any.whl

splifft/__init__.py

@@ -0,0 +1,12 @@
+"""Lightweight utilities for music source separation."""
+
+from pathlib import Path
+
+# TODO: this should be in the io module and we should use user cache paths to store models
+PATH_MODULE = Path(__file__).parent.parent
+PATH_BASE = PATH_MODULE.parent
+PATH_DATA = PATH_BASE / "data"
+PATH_CONFIG = PATH_DATA / "config"
+PATH_MODELS = PATH_DATA / "models"
+
+# NOTE: not re-exporting because our structure is simple enough.

splifft/__main__.py

@@ -0,0 +1,197 @@
+import logging
+from pathlib import Path
+from typing import Annotated, Optional
+
+import typer
+from rich.logging import RichHandler
+
+logging.basicConfig(level="INFO", format="%(message)s", datefmt="[%X]", handlers=[RichHandler()])
+logger = logging.getLogger(__name__)
+
+app = typer.Typer(
+    pretty_exceptions_show_locals=False,
+    help="A CLI for source separation.",
+    no_args_is_help=True,
+)
+
+
+# TODO: migrate away from hardcoding.
+_DEFAULT_MODULE_NAME = "splifft.models.bs_roformer"
+_DEFAULT_CLASS_NAME = "BSRoformer"
+
+
+@app.command()
+def separate(
+    mixture_path: Annotated[
+        Path,
+        typer.Argument(
+            ...,
+            exists=True,
+            file_okay=True,
+            dir_okay=True,
+            readable=True,
+            help="Path to the audio file to be separated.",
+        ),
+    ],
+    config_path: Annotated[
+        Path,
+        typer.Option(
+            "--config",
+            exists=True,
+            file_okay=True,
+            dir_okay=False,
+            readable=True,
+            help="Path to the model's JSON configuration file.",
+        ),
+    ],
+    checkpoint_path: Annotated[
+        Path,
+        typer.Option(
+            "--checkpoint",
+            exists=True,
+            file_okay=True,
+            dir_okay=False,
+            readable=True,
+            help="Path to the model's `.pt` or `.ckpt` checkpoint file.",
+        ),
+    ],
+    module_name: Annotated[
+        str,
+        typer.Option(
+            "--module",
+            help="Python module containing the model and configuration class.",
+        ),
+    ] = _DEFAULT_MODULE_NAME,
+    class_name: Annotated[
+        str,
+        typer.Option(
+            "--class",
+            help="Name of the model class to load from the module.",
+        ),
+    ] = _DEFAULT_CLASS_NAME,
+    package_name: Annotated[
+        Optional[str],
+        typer.Option(
+            "--package",
+            help=(
+                "The package to use as the anchor point from which to resolve the relative import "
+                "to an absolute import. This is only required when performing a relative import."
+            ),
+        ),
+    ] = None,
+    output_dir: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--output",
+            "-o",
+            file_okay=False,
+            dir_okay=True,
+            writable=True,
+            help="Directory to save the separated audio stems.",
+        ),
+    ] = None,
+    cpu: Annotated[
+        bool, typer.Option("--cpu", help="Force processing on CPU, even if CUDA is available.")
+    ] = False,
+) -> None:
+    """Separates an audio file into its constituent stems."""
+    import torch
+    import torchaudio
+
+    from .config import Config
+    from .core import get_dtype
+    from .inference import run_inference_on_file
+    from .io import load_weights, read_audio
+    from .models import ModelMetadata
+
+    use_cuda = not cpu and torch.cuda.is_available()
+    device = torch.device("cuda" if use_cuda else "cpu")
+    logger.info(f"using {device=}")
+
+    logger.info(f"loading configuration from {config_path=}")
+    config = Config.from_file(config_path)
+
+    logger.info(f"loading model metadata `{class_name}` from module `{module_name}`")
+    model_metadata = ModelMetadata.from_module(
+        module_name=module_name,
+        model_cls_name=class_name,
+        model_type=config.model_type,
+        package=package_name,
+    )
+    config_model_concrete = config.model.to_concrete(model_metadata.config)
+    model = model_metadata.model(config_model_concrete)
+    if config.inference.force_weights_dtype:
+        model = model.to(get_dtype(config.inference.force_weights_dtype))
+    logger.info(f"loading weights from {checkpoint_path=}")
+    model = load_weights(model, checkpoint_path, device).eval()
+    if (c := config.inference.compile_model) is not None:
+        logger.info("enabled torch compilation")
+        model = torch.compile(model, fullgraph=c.fullgraph, dynamic=c.dynamic, mode=c.mode)  # type: ignore
+
+    mixture_paths = mixture_path.glob("*") if mixture_path.is_dir() else [mixture_path]
+    for mixture_path in mixture_paths:
+        logger.info(f"processing audio file: {mixture_path=}")
+        mixture = read_audio(
+            mixture_path,
+            config.audio_io.target_sample_rate,
+            config.audio_io.force_channels,
+            device=device,
+        )
+        output_stems = run_inference_on_file(
+            mixture,
+            config=config,
+            model=model,
+        )
+        if not config.output:
+            return
+        curr_output_dir = output_dir or Path("./data/audio/output") / mixture_path.stem
+        curr_output_dir.mkdir(parents=True, exist_ok=True)
+        for stem_name, stem_data in output_stems.items():
+            if config.output.stem_names != "all" and stem_name not in config.output.stem_names:
+                continue
+
+            output_file = (curr_output_dir / stem_name).with_suffix(f".{config.output.file_format}")
+            torchaudio.save(
+                output_file,
+                stem_data.cpu(),
+                mixture.sample_rate,
+                format=config.output.file_format,
+                encoding=config.output.audio_encoding,
+                bits_per_sample=config.output.bit_depth,
+                # TODO: compression, backend
+            )
+            logger.info(f"wrote stem `{stem_name}` to {output_file}")
+
+
+@app.command()
+def debug() -> None:
+    """Prints detailed information about the environment, dependencies, and hardware
+    for debugging purposes."""
+    import sys
+
+    logger.info(f"{sys.version=}")
+    logger.info(f"{sys.executable=}")
+    logger.info(f"{sys.platform=}")
+    import platform
+
+    logger.info(f"{platform.system()=} ({platform.release()})")
+    logger.info(f"{platform.machine()=}")
+    import torch
+
+    logger.info(f"{torch.__version__=}")
+    logger.info(f"{torch.cuda.is_available()=}")
+    if torch.cuda.is_available():
+        logger.info(f"{torch.cuda.device_count()=}")
+        logger.info(f"{torch.cuda.current_device()=}")
+        device = torch.cuda.current_device()
+        logger.info(f"{torch.cuda.get_device_name(device)=}")
+        logger.info(f"{torch.cuda.get_device_capability(device)=}")
+        logger.info(f"{torch.cuda.get_device_properties(device)=}")
+    import torchaudio
+
+    logger.info(f"{torchaudio.__version__=}")
+    logger.info(f"{torchaudio.list_audio_backends()=}")
+
+
+if __name__ == "__main__":
+    app()

splifft/config.py

@@ -0,0 +1,265 @@
+"""Configuration"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import (
+    Annotated,
+    Any,
+    Hashable,
+    Literal,
+    Sequence,
+    TypeAlias,
+    TypeVar,
+    Union,
+    get_args,
+)
+
+from annotated_types import Len
+from pydantic import (
+    AfterValidator,
+    BaseModel,
+    BeforeValidator,
+    ConfigDict,
+    Discriminator,
+    Field,
+    StringConstraints,
+    TypeAdapter,
+    model_validator,
+)
+from pydantic_core import PydanticCustomError
+from typing_extensions import Self
+
+from .core import (
+    AudioEncoding,
+    BatchSize,
+    BitDepth,
+    Channels,
+    ChunkSize,
+    Dtype,
+    FileFormat,
+    OverlapRatio,
+    PaddingMode,
+    SampleRate,
+    WindowShape,
+)
+from .models import ModelConfigLikeT, ModelType
+from .models import ModelOutputStemName as _ModelOutputStemName
+
+# NOTE: we are not using typing.TYPE_CHECKING because pydantic relies on that
+
+_PYDANTIC_STRICT_CONFIG = ConfigDict(strict=True, extra="forbid")
+
+_Item = TypeVar("_Item", bound=Hashable)
+
+
+def _to_tuple(sequence: Sequence[_Item]) -> tuple[_Item, ...]:
+    # this is so json arrays are converted to tuples
+    return tuple(sequence)
+
+
+Tuple = Annotated[tuple[_Item, ...], BeforeValidator(_to_tuple)]
+
+
+def _validate_unique_sequence(sequence: Sequence[_Item]) -> Sequence[_Item]:
+    # e.g. to ensure there are no duplicate stem names
+    if len(sequence) != len(set(sequence)):
+        raise PydanticCustomError("unique_sequence", "Sequence must contain unique items")
+    return sequence
+
+
+_S = TypeVar("_S")
+NonEmptyUnique = Annotated[
+    _S,
+    Len(min_length=1),
+    AfterValidator(_validate_unique_sequence),
+    Field(json_schema_extra={"unique_items": True}),
+]
+
+ModelInputStemName: TypeAlias = Literal["mixture"]
+ModelOutputStemName: TypeAlias = Annotated[_ModelOutputStemName, StringConstraints(min_length=1)]
+_INPUT_STEM_NAMES = get_args(ModelInputStemName)
+
+
+# NOTE: the ideal case would be to use an ADT whose variants are known at "compile" time, e.g. in Rust:
+# enum ModelConfig {
+#     BsRoformer { param_x: ..., param_y: ... },
+#     Demucs { param_x: ..., params_z: ... },
+# }
+# but downstream users may want to register their own models with different configurations,
+# so a discriminated enum wouldn't work here.
+# so, we effectively let Config.model_config be dyn ModelConfigLike (i.e. dict[str, Any])
+# and defer the validation of the model configuration until it is actually needed instead of doing it eagerly.
+class LazyModelConfig(BaseModel):
+    """A lazily validated model configuration.
+
+    Note that it is not guaranteed to be fully valid until `to_concrete` is called.
+    """
+
+    chunk_size: ChunkSize
+    output_stem_names: NonEmptyUnique[Tuple[ModelOutputStemName]]
+
+    def to_concrete(
+        self,
+        model_config: type[ModelConfigLikeT],
+        *,
+        pydantic_config: ConfigDict = ConfigDict(extra="forbid"),
+    ) -> ModelConfigLikeT:
+        """Validate against a real model configuration and convert to it.
+
+        :raises pydantic.ValidationError: if extra fields are present in the model configuration
+            that doesn't exist in the concrete model configuration.
+        """
+        model_config_concrete: ModelConfigLikeT = TypeAdapter(
+            type(
+                f"{model_config.__name__}Validator",
+                (model_config,),
+                {"__pydantic_config__": pydantic_config},
+            )  # needed for https://docs.pydantic.dev/latest/errors/usage_errors/#type-adapter-config-unused
+        ).validate_python(self.model_dump())  # type: ignore
+        return model_config_concrete
+
+    @property
+    def stem_names(self) -> tuple[ModelInputStemName | ModelOutputStemName, ...]:
+        """Returns the model's input and output stem names."""
+        return (*_INPUT_STEM_NAMES, *self.output_stem_names)
+
+    model_config = ConfigDict(
+        strict=True, extra="allow"
+    )  # extra fields are not validated until `to_concrete`
+
+
+class AudioIOConfig(BaseModel):
+    target_sample_rate: SampleRate = 44100
+    force_channels: Channels | None = 2
+    """Whether to force mono or stereo audio input. If None, keep original."""
+
+    model_config = _PYDANTIC_STRICT_CONFIG
+
+
+class TorchCompileConfig(BaseModel):
+    fullgraph: bool = True
+    dynamic: bool = True
+    mode: Literal["default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs"] = (
+        "reduce-overhead"
+    )
+
+
+class InferenceConfig(BaseModel):
+    normalize_input_audio: bool = False
+    batch_size: BatchSize = 8
+    force_weights_dtype: Dtype | None = None
+    use_autocast_dtype: Dtype | None = None
+    compile_model: TorchCompileConfig | None = None
+    apply_tta: bool = False
+
+    model_config = _PYDANTIC_STRICT_CONFIG
+
+
+class ChunkingConfig(BaseModel):
+    method: Literal["overlap_add_windowed"] = "overlap_add_windowed"
+    overlap_ratio: OverlapRatio = 0.5
+    window_shape: WindowShape = "hann"
+    padding_mode: PaddingMode = "reflect"
+
+    model_config = _PYDANTIC_STRICT_CONFIG
+
+
+DerivedStemName: TypeAlias = Annotated[str, StringConstraints(min_length=1)]
+"""The name of a derived stem, e.g. `vocals_minus_drums`."""
+StemName: TypeAlias = Union[ModelOutputStemName, DerivedStemName]
+"""A name of a stem, either a model output stem or a derived stem."""
+
+
+class SubtractConfig(BaseModel):
+    operation: Literal["subtract"]
+    stem_name: StemName
+    by_stem_name: StemName
+
+    model_config = _PYDANTIC_STRICT_CONFIG
+
+
+class SumConfig(BaseModel):
+    operation: Literal["sum"]
+    stem_names: NonEmptyUnique[Tuple[StemName]]
+
+    model_config = _PYDANTIC_STRICT_CONFIG
+
+
+DerivedStemRule: TypeAlias = Annotated[Union[SubtractConfig, SumConfig], Discriminator("operation")]
+DerivedStemsConfig: TypeAlias = dict[DerivedStemName, DerivedStemRule]
+
+
+class OutputConfig(BaseModel):
+    stem_names: Literal["all"] | NonEmptyUnique[Tuple[StemName]] = "all"
+    file_format: FileFormat = "wav"
+    audio_encoding: AudioEncoding = "PCM_F"
+    bit_depth: BitDepth = 32
+
+    model_config = _PYDANTIC_STRICT_CONFIG
+
+
+# if we were to implement a model registry (which we shouldn't need)
+# heavily consider https://peps.python.org/pep-0487/#subclass-registration
+
+
+class Config(BaseModel):
+    identifier: str
+    """Unique identifier for this configuration"""
+    model_type: ModelType
+    model: LazyModelConfig
+    audio_io: AudioIOConfig = Field(default_factory=AudioIOConfig)
+    inference: InferenceConfig = Field(default_factory=InferenceConfig)
+    chunking: ChunkingConfig = Field(default_factory=ChunkingConfig)
+    derived_stems: DerivedStemsConfig | None = None
+    output: OutputConfig = Field(default_factory=OutputConfig)
+    experimental: dict[str, Any] | None = None
+    """Any extra experimental configurations outside of the `splifft` core."""
+
+    @model_validator(mode="after")
+    def check_derived_stems(self) -> Self:
+        if self.derived_stems is None:
+            return self
+        # accumulate valid stem names
+        existing_stem_names: list[StemName] = list(self.model.stem_names)
+        for derived_stem_name, definition in self.derived_stems.items():
+            if derived_stem_name in existing_stem_names:
+                raise PydanticCustomError(
+                    "derived_stem_name_conflict",
+                    "Derived stem `{derived_stem_name}` must not conflict with existing stem names: `{existing_stem_names}`",
+                    {
+                        "derived_stem_name": derived_stem_name,
+                        "existing_stem_names": existing_stem_names,
+                    },
+                )
+            required_stems: tuple[StemName, ...] = tuple()
+            if isinstance(definition, SubtractConfig):
+                required_stems = (definition.stem_name, definition.by_stem_name)
+            elif isinstance(definition, SumConfig):
+                required_stems = definition.stem_names
+            for stem_name in required_stems:
+                if stem_name not in existing_stem_names:
+                    raise PydanticCustomError(
+                        "invalid_derived_stem",
+                        "Derived stem `{derived_stem_name}` requires stem `{stem_name}` but is not found in `{existing_stem_names}`",
+                        {
+                            "derived_stem_name": derived_stem_name,
+                            "stem_name": stem_name,
+                            "existing_stem_names": existing_stem_names,
+                        },
+                    )
+            existing_stem_names.append(derived_stem_name)
+        return self
+
+    @classmethod
+    def from_file(cls, path: Path) -> Config:
+        with open(path, "r") as f:
+            config_data = json.load(f)
+        return Config.model_validate(config_data)
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,  # for .model
+        strict=True,
+        extra="forbid",
+    )