60 lines
2.7 KiB
Python
60 lines
2.7 KiB
Python
"""Offline, bounded HTDemucs inference. Model installation is a separate operator step."""
|
|
import json
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
for key in ('OMP_NUM_THREADS', 'OPENBLAS_NUM_THREADS', 'MKL_NUM_THREADS'):
|
|
os.environ[key] = '2'
|
|
CHECKSUM = "8726e21a993978c7ba086d3872e7608d7d5bfca646ca4aca459ffda844faa8b4"
|
|
SIGNATURE = '955717e8'
|
|
|
|
def available():
|
|
import importlib.util
|
|
root = Path(os.environ.get('SMART_EDIT_SEPARATION_MODEL_DIR', '/nonexistent'))
|
|
return importlib.util.find_spec('demucs') is not None and any(root.glob(SIGNATURE + '-*.th'))
|
|
|
|
def separate(source, destination):
|
|
if not available():
|
|
raise RuntimeError('人声分离模型尚未配置，请安装 Demucs 依赖和官方 HTDemucs 权重。')
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import torch
|
|
from demucs.states import load_model
|
|
import hashlib
|
|
from demucs.apply import apply_model
|
|
torch.set_num_threads(2)
|
|
torch.manual_seed(0)
|
|
audio, sr = sf.read(source, dtype='float32', always_2d=True)
|
|
if sr != 44100 or audio.shape[1] != 2 or not 0 < len(audio) <= sr * 180:
|
|
raise ValueError('分离输入必须为 180 秒以内的 44.1kHz 双声道音频')
|
|
if not np.isfinite(audio).all():
|
|
raise ValueError('音频包含无效采样')
|
|
checkpoint = next(Path(os.environ['SMART_EDIT_SEPARATION_MODEL_DIR']).glob(SIGNATURE + '-*.th'))
|
|
if hashlib.sha256(checkpoint.read_bytes()).hexdigest() != CHECKSUM:
|
|
raise ValueError('人声分离权重校验失败')
|
|
# Official checkpoint verified above; Demucs 4 uses a pickled model class.
|
|
model = load_model(torch.load(checkpoint, map_location='cpu', weights_only=False))
|
|
model.eval()
|
|
wave = torch.from_numpy(audio.T.copy())
|
|
reference = wave.mean(0)
|
|
mean, std = reference.mean(), reference.std()
|
|
if std < 1e-8:
|
|
vocals = np.zeros_like(audio)
|
|
else:
|
|
with torch.inference_mode():
|
|
output = apply_model(model, ((wave - mean) / std)[None], device='cpu',
|
|
shifts=1, split=True, overlap=.25, segment=7,
|
|
progress=False, num_workers=0)[0]
|
|
vocals = (output[model.sources.index('vocals')] * std + mean).T.numpy()
|
|
# Complementary stems reconstruct the original exactly, with no independent gain normalization.
|
|
destination = Path(destination)
|
|
destination.mkdir(parents=True, exist_ok=True)
|
|
sf.write(destination / 'vocals.wav', vocals, sr, subtype='FLOAT')
|
|
sf.write(destination / 'background.wav', audio - vocals, sr, subtype='FLOAT')
|
|
|
|
if __name__ == '__main__':
|
|
if sys.argv[1:] == ['--check']:
|
|
print(json.dumps({'separation': available()}))
|
|
else:
|
|
separate(*sys.argv[1:3])
|