audio-super-resolution 0.1.0__py3-none-any.whl

audio_super_resolution/__init__.py

@@ -0,0 +1,142 @@
+"""Audio super-resolution tools for Python."""
+
+from .audiosr_backend import AUDIOSR_MODEL_NAMES, AUDIOSR_SAMPLE_RATE, AudiosrBackend
+from .backends import register_backend, registered_backend_types
+from .config import InferenceConfig, default_model_cache_dir
+from .devices import DeviceInfo, available_devices, resolve_device
+from .downloads import (
+    download_model_weights as download_model_weights_for_spec,
+)
+from .downloads import (
+    download_weights_for_spec,
+    register_weight_provider,
+)
+from .manifest import (
+    ManifestComparison,
+    ManifestDifference,
+    build_manifest,
+    compare_manifests,
+    format_manifest_comparison,
+    load_manifest,
+    manifest_comparison_to_dict,
+    write_manifest,
+)
+from .model_weights import (
+    download_model_weights,
+    resolve_backend_model_weights,
+    resolve_model_weights,
+    verify_model_weights,
+)
+from .models import ModelInfo, find_model_spec, get_model_spec, list_models
+from .preprocess import DEFAULT_LOWPASS_CUTOFF_HZ, apply_preprocessing, lowpass_filter
+from .quality import (
+    AudioQualityReport,
+    build_quality_report_bundle,
+    format_quality_report,
+    inspect_audio_quality,
+    quality_report_to_dict,
+    write_quality_report_bundle,
+)
+from .resolver import (
+    DEFAULT_AUDIO_EXTENSIONS,
+    AudioSuperResolver,
+    BackendInfo,
+    EnhancementResult,
+    PlannedEnhancement,
+    available_backends,
+    discover_audio_files,
+    get_backend,
+    plan_enhancements,
+)
+from .specs import BackendCapability, ModelSpec, WeightFileSpec
+from .weight_store import (
+    ResolvedWeights,
+    resolve_weights_for_spec,
+    validate_weight_manifest_matches_spec,
+    verify_weights_for_spec,
+)
+from .weights import (
+    WeightFile,
+    WeightManifest,
+    load_safetensors,
+    read_weight_manifest,
+    resolve_manifest_file_paths,
+    resolve_weight_file_path,
+    resolve_weight_path,
+    sha256_file,
+    validate_weight_file_path,
+    verify_weight_file,
+    verify_weight_manifest,
+    write_weight_manifest,
+)
+
+__all__ = [
+    "AUDIOSR_MODEL_NAMES",
+    "AUDIOSR_SAMPLE_RATE",
+    "DEFAULT_AUDIO_EXTENSIONS",
+    "DEFAULT_LOWPASS_CUTOFF_HZ",
+    "AudioQualityReport",
+    "AudioSuperResolver",
+    "AudiosrBackend",
+    "BackendCapability",
+    "BackendInfo",
+    "DeviceInfo",
+    "EnhancementResult",
+    "InferenceConfig",
+    "ManifestComparison",
+    "ManifestDifference",
+    "ModelInfo",
+    "ModelSpec",
+    "PlannedEnhancement",
+    "ResolvedWeights",
+    "WeightFile",
+    "WeightFileSpec",
+    "WeightManifest",
+    "apply_preprocessing",
+    "available_backends",
+    "available_devices",
+    "build_manifest",
+    "build_quality_report_bundle",
+    "compare_manifests",
+    "default_model_cache_dir",
+    "discover_audio_files",
+    "download_model_weights",
+    "download_model_weights_for_spec",
+    "download_weights_for_spec",
+    "find_model_spec",
+    "format_manifest_comparison",
+    "format_quality_report",
+    "get_backend",
+    "get_model_spec",
+    "inspect_audio_quality",
+    "list_models",
+    "load_manifest",
+    "load_safetensors",
+    "lowpass_filter",
+    "manifest_comparison_to_dict",
+    "plan_enhancements",
+    "quality_report_to_dict",
+    "read_weight_manifest",
+    "register_backend",
+    "register_weight_provider",
+    "registered_backend_types",
+    "resolve_backend_model_weights",
+    "resolve_device",
+    "resolve_manifest_file_paths",
+    "resolve_model_weights",
+    "resolve_weights_for_spec",
+    "resolve_weight_file_path",
+    "resolve_weight_path",
+    "sha256_file",
+    "validate_weight_file_path",
+    "validate_weight_manifest_matches_spec",
+    "verify_model_weights",
+    "verify_weights_for_spec",
+    "verify_weight_file",
+    "verify_weight_manifest",
+    "write_manifest",
+    "write_quality_report_bundle",
+    "write_weight_manifest",
+]
+
+__version__ = "0.1.0"

audio_super_resolution/audiosr_backend.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from .backends.audiosr_external import (
+    AUDIOSR_MODEL_NAMES,
+    AUDIOSR_SAMPLE_RATE,
+    AudiosrBackend,
+    _import_audiosr,
+    _waveform_to_audio_array,
+)
+
+__all__ = [
+    "AUDIOSR_MODEL_NAMES",
+    "AUDIOSR_SAMPLE_RATE",
+    "AudiosrBackend",
+    "_import_audiosr",
+    "_waveform_to_audio_array",
+]

audio_super_resolution/backends/__init__.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from .audiosr_external import AUDIOSR_MODEL_NAMES, AUDIOSR_SAMPLE_RATE, AudiosrBackend
+from .base import BackendInfo, EnhancementBackend, FileEnhancementBackend
+from .lavasr_compat import LAVASR_MODEL_ID, LAVASR_REVISION, LAVASR_SAMPLE_RATE, LavaSRCompatBackend
+from .registry import (
+    available_backends,
+    get_backend,
+    list_backend_model_specs,
+    register_backend,
+    registered_backend_types,
+)
+from .sinc import SincResampleBackend
+
+__all__ = [
+    "AUDIOSR_MODEL_NAMES",
+    "AUDIOSR_SAMPLE_RATE",
+    "LAVASR_MODEL_ID",
+    "LAVASR_REVISION",
+    "LAVASR_SAMPLE_RATE",
+    "AudiosrBackend",
+    "BackendInfo",
+    "EnhancementBackend",
+    "FileEnhancementBackend",
+    "LavaSRCompatBackend",
+    "SincResampleBackend",
+    "available_backends",
+    "get_backend",
+    "list_backend_model_specs",
+    "register_backend",
+    "registered_backend_types",
+]

audio_super_resolution/backends/audiosr_external.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import importlib
+import importlib.util
+import os
+from pathlib import Path
+from types import ModuleType
+
+import numpy as np
+import soundfile as sf
+
+from ..config import InferenceConfig
+from ..specs import ModelSpec
+from .base import DEFAULT_FILE_BACKEND_CAPABILITY
+
+AUDIOSR_SAMPLE_RATE = 48000
+AUDIOSR_MODEL_NAMES = ("basic", "speech")
+
+
+class AudiosrBackend:
+    """AudioSR latent diffusion backend.
+
+    The heavy `audiosr` dependency is imported only when this backend is selected.
+    """
+
+    name = "audiosr"
+    description = "AudioSR latent diffusion backend. Requires the optional audiosr dependency."
+    optional_dependency = "audiosr"
+    package_extra = "audiosr"
+
+    def __init__(self, config: InferenceConfig | None = None) -> None:
+        self.config = config or InferenceConfig()
+        self._model = None
+
+    @classmethod
+    def is_available(cls) -> bool:
+        try:
+            return importlib.util.find_spec("audiosr") is not None
+        except (ImportError, ValueError):
+            return False
+
+    @classmethod
+    def model_specs(cls) -> tuple[ModelSpec, ...]:
+        return tuple(
+            ModelSpec(
+                id=f"audiosr-{model_name}",
+                backend=cls.name,
+                model_name=model_name,
+                name=f"AudioSR {model_name.title()}",
+                description=f"AudioSR latent diffusion {model_name} model.",
+                implementation="external_package",
+                domain=("speech", "music", "sfx", "general"),
+                architecture="latent-diffusion-audio-super-resolution",
+                target_sample_rates=(AUDIOSR_SAMPLE_RATE,),
+                weights_license="unknown",
+                weights_source="Hugging Face download managed by the audiosr package",
+                maturity="experimental",
+                capability=DEFAULT_FILE_BACKEND_CAPABILITY,
+            )
+            for model_name in AUDIOSR_MODEL_NAMES
+        )
+
+    def enhance(self, audio: np.ndarray, sample_rate: int, target_sample_rate: int) -> np.ndarray:
+        raise RuntimeError("The audiosr backend requires file-based enhancement through AudioSuperResolver.enhance().")
+
+    def enhance_file(self, input_path: str | Path, output_path: str | Path, target_sample_rate: int) -> None:
+        if target_sample_rate != AUDIOSR_SAMPLE_RATE:
+            raise ValueError("The audiosr backend outputs 48000 Hz audio; set target_sr=48000.")
+        if self.config.device == "directml":
+            raise ValueError("The audiosr backend does not support directml; use cpu, cuda, mps, or auto.")
+        if self.config.precision not in {"float32", "auto"}:
+            raise ValueError("The audiosr backend currently supports only float32 or auto precision.")
+        if self.config.model_name not in AUDIOSR_MODEL_NAMES:
+            choices = ", ".join(AUDIOSR_MODEL_NAMES)
+            raise ValueError(f"The audiosr backend model_name must be one of: {choices}")
+
+        audiosr = _import_audiosr()
+        self._prepare_cache_environment()
+
+        model = self._load_model(audiosr)
+        waveform = audiosr.super_resolution(
+            model,
+            str(input_path),
+            seed=self.config.seed,
+            guidance_scale=self.config.guidance_scale,
+            ddim_steps=self.config.ddim_steps,
+        )
+
+        output_path = Path(output_path)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        sf.write(output_path, _waveform_to_audio_array(waveform), samplerate=AUDIOSR_SAMPLE_RATE)
+
+    def _prepare_cache_environment(self) -> None:
+        cache_dir = self.config.ensure_model_cache_dir()
+        os.environ.setdefault("HF_HOME", str(cache_dir / "huggingface"))
+
+    def _load_model(self, audiosr: ModuleType):
+        if self._model is None:
+            self._model = audiosr.build_model(model_name=self.config.model_name, device=self.config.device)
+        return self._model
+
+
+def _import_audiosr() -> ModuleType:
+    try:
+        return importlib.import_module("audiosr")
+    except ImportError as exc:
+        raise RuntimeError(
+            "The audiosr backend requires the optional audiosr dependency. "
+            "Install it with `pip install audio-super-resolution[audiosr]`."
+        ) from exc
+
+
+def _waveform_to_audio_array(waveform) -> np.ndarray:
+    if hasattr(waveform, "detach"):
+        waveform = waveform.detach().cpu().numpy()
+
+    audio = np.asarray(waveform, dtype=np.float32)
+
+    if audio.ndim == 3:
+        audio = audio[0]
+    if audio.ndim == 2 and audio.shape[0] <= 8:
+        audio = audio.T
+    if audio.ndim == 2 and audio.shape[1] == 1:
+        audio = audio[:, 0]
+    if audio.ndim not in {1, 2}:
+        raise ValueError(f"Unsupported AudioSR waveform shape: {audio.shape}")
+
+    return audio

audio_super_resolution/backends/base.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Protocol
+
+import numpy as np
+
+from ..config import InferenceConfig
+from ..specs import BackendCapability, ModelSpec
+
+
+@dataclass(frozen=True)
+class BackendInfo:
+    """User-facing metadata for an enhancement backend."""
+
+    name: str
+    description: str
+    installed: bool
+    optional_dependency: str | None = None
+    package_extra: str | None = None
+
+
+class EnhancementBackend(Protocol):
+    """Interface implemented by audio super-resolution backends."""
+
+    name: str
+    description: str
+    config: InferenceConfig
+
+    def enhance(self, audio: np.ndarray, sample_rate: int, target_sample_rate: int) -> np.ndarray:
+        """Return audio enhanced to target_sample_rate."""
+
+
+class FileEnhancementBackend(EnhancementBackend, Protocol):
+    """Optional file-native backend interface."""
+
+    def enhance_file(self, input_path: str | Path, output_path: str | Path, target_sample_rate: int) -> None:
+        """Write enhanced audio to output_path."""
+
+
+DEFAULT_ARRAY_BACKEND_CAPABILITY = BackendCapability(
+    supports_array_io=True,
+    supports_file_io=False,
+    supports_chunking=True,
+    deterministic=True,
+    supports_cpu=True,
+)
+
+
+DEFAULT_FILE_BACKEND_CAPABILITY = BackendCapability(
+    supports_array_io=False,
+    supports_file_io=True,
+    supports_chunking=True,
+    deterministic=False,
+    supports_cpu=True,
+    supports_cuda=True,
+    supports_mps=True,
+)
+
+
+def backend_model_specs(backend_type: type[EnhancementBackend]) -> tuple[ModelSpec, ...]:
+    """Return model specs exposed by a backend type."""
+
+    model_specs = getattr(backend_type, "model_specs", None)
+    if callable(model_specs):
+        return tuple(model_specs())
+    return ()

audio_super_resolution/backends/lavasr_compat.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+import importlib.util
+
+import numpy as np
+
+from ..config import InferenceConfig
+from ..specs import BackendCapability, ModelSpec, WeightFileSpec
+
+LAVASR_SAMPLE_RATE = 48000
+LAVASR_MODEL_ID = "lavasr-v2-bwe"
+LAVASR_REVISION = "b98dc8be472da45ab7b6346ad7997e1dfeb5911d"
+
+
+class LavaSRCompatBackend:
+    """Self-contained LavaSR-compatible backend placeholder for verified BWE weights."""
+
+    name = "lavasr-compat"
+    description = "Self-contained LavaSR-compatible speech bandwidth extension backend."
+    optional_dependency = "torch"
+    package_extra = "lavasr"
+
+    def __init__(self, config: InferenceConfig | None = None) -> None:
+        self.config = config or InferenceConfig()
+
+    @classmethod
+    def is_available(cls) -> bool:
+        return importlib.util.find_spec("torch") is not None and importlib.util.find_spec("yaml") is not None
+
+    @classmethod
+    def model_specs(cls) -> tuple[ModelSpec, ...]:
+        return (
+            ModelSpec(
+                id=LAVASR_MODEL_ID,
+                backend=cls.name,
+                model_name=LAVASR_MODEL_ID,
+                name="LavaSR v2 BWE",
+                description="LavaSR v2 bandwidth extension weights for 48 kHz output.",
+                implementation="self_torch",
+                domain=("speech",),
+                architecture="vocos-istft-bwe",
+                target_sample_rates=(LAVASR_SAMPLE_RATE,),
+                weights_license="Apache-2.0",
+                weights_source="YatharthS/LavaSR",
+                weight_provider="huggingface",
+                default_weight_revision=LAVASR_REVISION,
+                requires_weights=True,
+                maturity="experimental",
+                weight_files=(
+                    WeightFileSpec(
+                        path="enhancer_v2/config.yaml",
+                        sha256="0d970f5cdc1913730c417b49e476bb09bb8b874583d113f71a7f10c8bb3a4b7d",
+                        size=526,
+                    ),
+                    WeightFileSpec(
+                        path="enhancer_v2/pytorch_model.bin",
+                        sha256="d100db961b2c125d77a52a12215c689e44cd9926a72f117513395b7e25e6de12",
+                        size=56316591,
+                    ),
+                ),
+                capability=BackendCapability(
+                    supports_array_io=True,
+                    supports_file_io=False,
+                    supports_chunking=True,
+                    deterministic=True,
+                    supports_cpu=True,
+                    supports_cuda=True,
+                    supports_mps=True,
+                    precision_modes=("float32", "auto"),
+                ),
+            ),
+        )
+
+    def enhance(self, audio: np.ndarray, sample_rate: int, target_sample_rate: int) -> np.ndarray:
+        if target_sample_rate != LAVASR_SAMPLE_RATE:
+            raise ValueError("The lavasr-compat backend outputs 48000 Hz audio; set target_sr=48000.")
+        if self.config.denoise:
+            raise ValueError("denoise is reserved but not supported by lavasr-compat yet")
+
+        from ..weight_store import resolve_weights_for_spec
+
+        resolve_weights_for_spec(self.model_specs()[0], self.config, allow_download=self.config.download_weights)
+        raise RuntimeError("lavasr-compat weight management is available, but inference is not implemented yet.")

audio_super_resolution/backends/registry.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from ..config import InferenceConfig
+from ..specs import ModelSpec
+from .audiosr_external import AudiosrBackend
+from .base import BackendInfo, EnhancementBackend, backend_model_specs
+from .lavasr_compat import LavaSRCompatBackend
+from .sinc import SincResampleBackend
+
+_BACKENDS: dict[str, type[EnhancementBackend]] = {}
+
+
+def register_backend(backend_type: type[EnhancementBackend], *, replace: bool = False) -> None:
+    """Register an enhancement backend type by its public name."""
+
+    name = getattr(backend_type, "name", None)
+    if not isinstance(name, str) or not name:
+        raise ValueError("backend_type must define a non-empty string name")
+    if name in _BACKENDS and not replace:
+        raise ValueError(f"Backend {name!r} is already registered")
+    _BACKENDS[name] = backend_type
+
+
+def registered_backend_types() -> dict[str, type[EnhancementBackend]]:
+    """Return a copy of the backend registry."""
+
+    return dict(_BACKENDS)
+
+
+def available_backends() -> list[BackendInfo]:
+    """Return the registered enhancement backends."""
+
+    return [
+        BackendInfo(
+            name=name,
+            description=backend.description,
+            installed=_backend_is_available(backend),
+            optional_dependency=getattr(backend, "optional_dependency", None),
+            package_extra=getattr(backend, "package_extra", None),
+        )
+        for name, backend in sorted(_BACKENDS.items(), key=lambda item: item[0])
+    ]
+
+
+def get_backend(name: str, config: InferenceConfig | None = None) -> EnhancementBackend:
+    """Create a backend by name."""
+
+    try:
+        backend_type = _BACKENDS[name]
+    except KeyError as exc:
+        choices = ", ".join(sorted(_BACKENDS))
+        raise ValueError(f"Unknown backend {name!r}. Available backends: {choices}") from exc
+
+    return backend_type(config=config)
+
+
+def list_backend_model_specs() -> list[ModelSpec]:
+    """Return model specs exposed by every registered backend."""
+
+    specs: list[ModelSpec] = []
+    for _, backend_type in sorted(_BACKENDS.items(), key=lambda item: item[0]):
+        specs.extend(backend_model_specs(backend_type))
+    return specs
+
+
+def _backend_is_available(backend_type: type[EnhancementBackend]) -> bool:
+    checker = getattr(backend_type, "is_available", None)
+    if checker is None:
+        return True
+    return bool(checker())
+
+
+register_backend(AudiosrBackend)
+register_backend(LavaSRCompatBackend)
+register_backend(SincResampleBackend)

audio_super_resolution/backends/sinc.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import numpy as np
+from scipy.signal import resample_poly
+
+from ..config import InferenceConfig
+from ..specs import ModelSpec
+from .base import DEFAULT_ARRAY_BACKEND_CAPABILITY
+
+
+class SincResampleBackend:
+    """Deterministic baseline backend using polyphase sinc resampling."""
+
+    name = "sinc-resample"
+    description = "Deterministic polyphase sinc resampling baseline."
+    optional_dependency = None
+    package_extra = None
+
+    def __init__(self, config: InferenceConfig | None = None) -> None:
+        self.config = config or InferenceConfig()
+
+    @classmethod
+    def model_specs(cls) -> tuple[ModelSpec, ...]:
+        return (
+            ModelSpec(
+                id="sinc-resample",
+                backend=cls.name,
+                name="Sinc Resample Baseline",
+                description=cls.description,
+                implementation="baseline",
+                domain=("general",),
+                architecture="polyphase-sinc-resampling",
+                target_sample_rates=None,
+                weights_license=None,
+                weights_source=None,
+                maturity="stable",
+                capability=DEFAULT_ARRAY_BACKEND_CAPABILITY,
+            ),
+        )
+
+    def enhance(self, audio: np.ndarray, sample_rate: int, target_sample_rate: int) -> np.ndarray:
+        if sample_rate == target_sample_rate:
+            return np.asarray(audio)
+
+        gcd = np.gcd(sample_rate, target_sample_rate)
+        up = target_sample_rate // gcd
+        down = sample_rate // gcd
+        return resample_poly(audio, up, down, axis=0)