File size: 3,560 Bytes
81ba775 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | 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()
|