src/backend.py

from __future__ import annotations

import gc
from pathlib import Path

import torch
from audio_separator.separator import Separator

from .downloads import CustomModelRecord
from .observability import log_level_number, normalize_log_level


class RegisteredModelSeparator(Separator):
    """Adapter that lets the public package load trusted MDXC model+YAML pairs."""

    def __init__(self, *args, custom_models: list[CustomModelRecord] | None = None, **kwargs):
        self._custom_models = {record.filename: record for record in (custom_models or [])}
        super().__init__(*args, **kwargs)

    def configure_cuda(self, ort_providers):
        self.logger.info("CUDA is available in Torch, configuring GPU inference")
        self.torch_device = torch.device("cuda")
        if "CUDAExecutionProvider" in ort_providers:
            self.onnx_execution_provider = ["CUDAExecutionProvider", "CPUExecutionProvider"]
        else:
            self.logger.warning("ONNX Runtime CUDA provider is unavailable; ONNX models will use CPU")
            self.onnx_execution_provider = ["CPUExecutionProvider"]

    def download_model_files(self, model_filename):
        record = self._custom_models.get(model_filename)
        if record is None:
            return super().download_model_files(model_filename)
        self.model_friendly_name = record.display_name
        self.model_is_uvr_vip = False
        return (
            record.filename,
            record.architecture,
            record.display_name,
            str(record.model_path),
            str(record.config_path),
        )


def create_separator(
    output_dir: Path,
    model_dir: Path,
    custom_models: list[CustomModelRecord],
    output_format: str,
    ensemble_algorithm: str,
    single_stem: str | None,
    pitch_shift: int,
    chunk_duration: int | None,
    output_bitrate: str | None = None,
    sample_rate: int = 44100,
    normalization_threshold: float = 0.9,
    amplification_threshold: float = 0.0,
    log_level: str = "INFO",
    job_id: str | None = None,
) -> RegisteredModelSeparator:
    level_name = normalize_log_level(log_level)
    job_text = str(job_id or "unknown")
    return RegisteredModelSeparator(
        log_level=log_level_number(level_name),
        log_formatter=(
            "%(asctime)s - %(levelname)s - %(module)s - "
            f"job={job_text} - %(message)s"
        ),
        model_file_dir=str(model_dir),
        output_dir=str(output_dir),
        output_format=output_format,
        output_bitrate=output_bitrate,
        sample_rate=int(sample_rate),
        normalization_threshold=float(normalization_threshold),
        amplification_threshold=float(amplification_threshold),
        output_single_stem=single_stem,
        use_soundfile=False,
        use_autocast=bool(torch.cuda.is_available()),
        chunk_duration=chunk_duration,
        ensemble_algorithm=ensemble_algorithm,
        mdxc_params={
            "segment_size": 256,
            "override_model_segment_size": False,
            "batch_size": 1,
            "overlap": 8,
            "pitch_shift": int(pitch_shift),
        },
        custom_models=custom_models,
    )


def release_accelerators(separator=None) -> None:
    try:
        if separator is not None and getattr(separator, "model_instance", None) is not None:
            separator.model_instance.clear_gpu_cache()
    except Exception:
        pass
    del separator
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()