generated from dellevin/template
1016 lines
39 KiB
Python
1016 lines
39 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
UVR 人声分离蓝图
|
||
基于 Ultimate Vocal Remover GUI 的音频源分离能力
|
||
支持 VR / MDX-Net / Demucs 三种架构
|
||
"""
|
||
import os
|
||
import sys
|
||
import io
|
||
import json
|
||
import uuid
|
||
import hashlib
|
||
import zipfile
|
||
import shutil
|
||
import threading
|
||
import tempfile
|
||
import traceback
|
||
from flask import Blueprint, render_template, request, jsonify, send_file
|
||
import numpy as np
|
||
|
||
try:
|
||
from config import BASE_DIR
|
||
except ImportError:
|
||
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
|
||
bp = Blueprint('uvr_sep', __name__, url_prefix='/uvr-sep')
|
||
|
||
AUDIO_EXTS = ('.wav', '.flac', '.ogg', '.mp3', '.aac', '.m4a', '.wma', '.aiff', '.opus')
|
||
CONFIG_PATH = os.path.join(BASE_DIR, 'config', 'uvr_sep_config.json')
|
||
|
||
# 架构类型常量（与 UVR 的 gui_data.constants 一致）
|
||
VR_ARCH = 'VR Arc'
|
||
MDX_ARCH = 'MDX-Net'
|
||
DEMUCS_ARCH = 'Demucs'
|
||
|
||
_tasks = {}
|
||
|
||
# ── 配置读写 ──────────────────────────────────────────────────────────────────
|
||
|
||
DEFAULT_PARAMS = {
|
||
'uvr_project_path': '', 'model_dir_mode': 'absolute',
|
||
'demucs_model_dir': '', 'vr_model_dir': '', 'mdx_model_dir': '',
|
||
'arch_type': 'Demucs',
|
||
'save_format': 'wav', 'wav_type': 'PCM_16', 'mp3_bitrate': '320k',
|
||
'is_gpu': True, 'device_set': 'Default',
|
||
'is_normalization': False,
|
||
'is_primary_stem_only': False, 'is_secondary_stem_only': False,
|
||
'demucs_stems': 'All Stems',
|
||
'demucs_segment': 'Default',
|
||
'mdx_segment_size': 'Default', 'mdx_overlap': 0.25, 'mdx_batch_size': 1,
|
||
'vr_window_size': 1024, 'vr_aggression': 5, 'vr_batch_size': 4,
|
||
'demucs_selected_model': '', 'mdx_selected_model': '', 'vr_selected_model': '',
|
||
}
|
||
|
||
|
||
def _load_config():
|
||
cfg = dict(DEFAULT_PARAMS)
|
||
if os.path.exists(CONFIG_PATH):
|
||
try:
|
||
with open(CONFIG_PATH, 'r', encoding='utf-8') as f:
|
||
cfg.update(json.load(f))
|
||
except Exception:
|
||
pass
|
||
return cfg
|
||
|
||
|
||
def _save_config(cfg):
|
||
os.makedirs(os.path.dirname(CONFIG_PATH), exist_ok=True)
|
||
with open(CONFIG_PATH, 'w', encoding='utf-8') as f:
|
||
json.dump(cfg, f, ensure_ascii=False, indent=2)
|
||
|
||
|
||
def _resolve_model_dir(raw_dir, cfg=None):
|
||
"""将模型目录路径解析为绝对路径"""
|
||
if not raw_dir:
|
||
return ''
|
||
if os.path.isabs(raw_dir):
|
||
return raw_dir
|
||
if cfg is None:
|
||
cfg = _load_config()
|
||
uvr_path = cfg.get('uvr_project_path', '')
|
||
mode = cfg.get('model_dir_mode', 'absolute')
|
||
if mode == 'relative' and uvr_path:
|
||
return os.path.normpath(os.path.join(uvr_path, raw_dir))
|
||
# absolute 模式下如果输入的是相对路径，也尝试拼接 uvr_project_path
|
||
if uvr_path:
|
||
resolved = os.path.normpath(os.path.join(uvr_path, raw_dir))
|
||
if os.path.isdir(resolved):
|
||
return resolved
|
||
return raw_dir
|
||
|
||
|
||
def _get_model_dir_for_arch(arch_type, explicit_dir=None, cfg=None):
|
||
"""获取指定架构的模型目录：优先显式路径，否则从配置自动推断"""
|
||
if explicit_dir:
|
||
return _resolve_model_dir(explicit_dir, cfg)
|
||
if cfg is None:
|
||
cfg = _load_config()
|
||
key_map = {'Demucs': 'demucs_model_dir', 'VR Arc': 'vr_model_dir', 'MDX-Net': 'mdx_model_dir'}
|
||
raw = cfg.get(key_map.get(arch_type, ''), '')
|
||
return _resolve_model_dir(raw, cfg)
|
||
|
||
|
||
def _get_uvr_paths():
|
||
"""计算 UVR 项目内的关键路径"""
|
||
cfg = _load_config()
|
||
uvr_path = cfg.get('uvr_project_path', '')
|
||
if not uvr_path:
|
||
return {}
|
||
models_dir = os.path.join(uvr_path, 'models')
|
||
return {
|
||
'uvr_path': uvr_path,
|
||
'models_dir': models_dir,
|
||
'vr_param_dir': os.path.join(uvr_path, 'lib_v5', 'vr_network', 'modelparams'),
|
||
'mdx_c_config_path': os.path.join(models_dir, 'MDX_Net_Models', 'model_data', 'mdx_c_configs'),
|
||
'mixer_path': os.path.join(uvr_path, 'lib_v5', 'mixer.ckpt'),
|
||
'denoiser_path': os.path.join(models_dir, 'VR_Models', 'UVR-DeNoise-Lite.pth'),
|
||
'deverb_path': os.path.join(models_dir, 'VR_Models', 'UVR-DeEcho-DeReverb.pth'),
|
||
}
|
||
|
||
|
||
# ── 延迟导入 UVR ──────────────────────────────────────────────────────────────
|
||
|
||
_uvr_imported = False
|
||
_uvr_error = None
|
||
# 导入后的 UVR 模块引用
|
||
_ModelParameters = None
|
||
_SeperateVR = _SeperateMDX = _SeperateMDXC = _SeperateDemucs = None
|
||
_secondary_stem = None
|
||
|
||
|
||
def _ensure_uvr_imports():
|
||
"""延迟导入 UVR 模块，仅在使用时加载"""
|
||
global _uvr_imported, _uvr_error
|
||
global _ModelParameters, _secondary_stem
|
||
global _SeperateVR, _SeperateMDX, _SeperateMDXC, _SeperateDemucs
|
||
|
||
if _uvr_imported:
|
||
return True
|
||
if _uvr_error:
|
||
raise ImportError(_uvr_error)
|
||
|
||
cfg = _load_config()
|
||
uvr_path = cfg.get('uvr_project_path', '')
|
||
if not uvr_path or not os.path.isdir(uvr_path):
|
||
_uvr_error = '请先配置 UVR 项目路径'
|
||
raise ImportError(_uvr_error)
|
||
|
||
try:
|
||
if uvr_path not in sys.path:
|
||
sys.path.insert(0, uvr_path)
|
||
|
||
from gui_data.constants import secondary_stem as _ss
|
||
_secondary_stem = _ss
|
||
|
||
from lib_v5.vr_network.model_param_init import ModelParameters as _MP
|
||
_ModelParameters = _MP
|
||
|
||
from separate import (
|
||
SeperateVR as _SVR, SeperateMDX as _SMDX,
|
||
SeperateMDXC as _SMDXC, SeperateDemucs as _SDem,
|
||
)
|
||
_SeperateVR = _SVR
|
||
_SeperateMDX = _SMDX
|
||
_SeperateMDXC = _SMDXC
|
||
_SeperateDemucs = _SDem
|
||
|
||
# PyTorch 2.6+ 默认 weights_only=True，对 UVR 模型不兼容
|
||
# patch torch.load，对 demucs 的模型加载使用 weights_only=False
|
||
import torch as _torch
|
||
_original_torch_load = _torch.load
|
||
def _safe_torch_load(*args, **kwargs):
|
||
if 'weights_only' not in kwargs:
|
||
kwargs['weights_only'] = False
|
||
return _original_torch_load(*args, **kwargs)
|
||
_torch.load = _safe_torch_load
|
||
|
||
# 重新检测 CUDA（separate.py 在模块加载时检测一次，可能不准确）
|
||
import separate as _separate_module
|
||
_separate_module.cuda_available = _torch.cuda.is_available()
|
||
|
||
# 修复 UVR demucs/apply.py 中 bag_num/prog_bar 未初始化的 bug
|
||
try:
|
||
from demucs import apply as _demucs_apply
|
||
_orig_apply = _demucs_apply.apply_model
|
||
def _patched_apply(*args, **kwargs):
|
||
_demucs_apply.bag_num = getattr(_demucs_apply, 'bag_num', 1)
|
||
_demucs_apply.prog_bar = getattr(_demucs_apply, 'prog_bar', 0)
|
||
return _orig_apply(*args, **kwargs)
|
||
_demucs_apply.apply_model = _patched_apply
|
||
except Exception:
|
||
pass
|
||
|
||
# 修复 librosa 新版本 API 不兼容（位置参数 → 关键字参数）
|
||
import librosa as _librosa
|
||
_orig_librosa_load = _librosa.load
|
||
def _compat_librosa_load(path, *args, **kwargs):
|
||
if args:
|
||
if 'sr' not in kwargs and len(args) >= 1:
|
||
kwargs['sr'] = args[0]
|
||
if 'mono' not in kwargs and len(args) >= 2:
|
||
kwargs['mono'] = args[1]
|
||
args = ()
|
||
return _orig_librosa_load(path, *args, **kwargs)
|
||
_librosa.load = _compat_librosa_load
|
||
|
||
_orig_librosa_stft = _librosa.stft
|
||
def _compat_librosa_stft(y, *args, **kwargs):
|
||
names = ['n_fft', 'hop_length', 'win_length', 'window', 'center',
|
||
'pad_mode', 'length', 'return_complex']
|
||
for i, v in enumerate(args):
|
||
if i < len(names) and names[i] not in kwargs:
|
||
kwargs[names[i]] = v
|
||
return _orig_librosa_stft(y, **kwargs)
|
||
_librosa.stft = _compat_librosa_stft
|
||
|
||
_orig_librosa_resample = _librosa.resample
|
||
def _compat_librosa_resample(y, *args, **kwargs):
|
||
names = ['orig_sr', 'target_sr', 'fix', 'scale', 'axis', 'res_type']
|
||
for i, v in enumerate(args):
|
||
if i < len(names) and names[i] not in kwargs:
|
||
kwargs[names[i]] = v
|
||
return _orig_librosa_resample(y, **kwargs)
|
||
_librosa.resample = _compat_librosa_resample
|
||
|
||
_orig_librosa_istft = getattr(_librosa, 'istft', None)
|
||
if _orig_librosa_istft:
|
||
def _compat_librosa_istft(stft_matrix, *args, **kwargs):
|
||
names = ['hop_length', 'win_length', 'window', 'center', 'length', 'dtype']
|
||
for i, v in enumerate(args):
|
||
if i < len(names) and names[i] not in kwargs:
|
||
kwargs[names[i]] = v
|
||
return _orig_librosa_istft(stft_matrix, **kwargs)
|
||
_librosa.istft = _compat_librosa_istft
|
||
|
||
# 修复 VR Arc: cmb_spectrogram_to_wave 中 np.ndarray(dtype=complex)
|
||
# 未初始化包含垃圾值 → 用 np.zeros 替代
|
||
from lib_v5 import spec_utils as _spec_utils
|
||
import numpy as _cmb_np
|
||
import librosa as _cmb_librosa
|
||
def _fixed_cmb_spectrogram_to_wave(spec_m, mp, extra_bins_h=None, extra_bins=None, is_v51_model=False):
|
||
spec_m = _cmb_np.nan_to_num(spec_m, nan=0.0, posinf=0.0, neginf=0.0)
|
||
bands_n = len(mp.param['band'])
|
||
offset = 0
|
||
for d in range(1, bands_n + 1):
|
||
bp = mp.param['band'][d]
|
||
spec_s = _cmb_np.zeros(shape=(2, bp['n_fft'] // 2 + 1, spec_m.shape[2]), dtype=complex)
|
||
h = bp['crop_stop'] - bp['crop_start']
|
||
spec_s[:, bp['crop_start']:bp['crop_stop'], :] = spec_m[:, offset:offset+h, :]
|
||
offset += h
|
||
if d == bands_n:
|
||
if extra_bins_h:
|
||
max_bin = bp['n_fft'] // 2
|
||
spec_s[:, max_bin-extra_bins_h:max_bin, :] = extra_bins[:, :extra_bins_h, :]
|
||
if bp['hpf_start'] > 0:
|
||
if is_v51_model:
|
||
spec_s *= _spec_utils.get_hp_filter_mask(spec_s.shape[1], bp['hpf_start'], bp['hpf_stop'] - 1)
|
||
else:
|
||
spec_s = _spec_utils.fft_hp_filter(spec_s, bp['hpf_start'], bp['hpf_stop'] - 1)
|
||
if bands_n == 1:
|
||
wav = _spec_utils.spectrogram_to_wave(spec_s, bp['hl'], mp, d, is_v51_model)
|
||
else:
|
||
wav = _cmb_np.add(wav, _spec_utils.spectrogram_to_wave(spec_s, bp['hl'], mp, d, is_v51_model))
|
||
else:
|
||
sr = mp.param['band'][d+1]['sr']
|
||
if d == 1:
|
||
if is_v51_model:
|
||
spec_s *= _spec_utils.get_lp_filter_mask(spec_s.shape[1], bp['lpf_start'], bp['lpf_stop'])
|
||
else:
|
||
spec_s = _spec_utils.fft_lp_filter(spec_s, bp['lpf_start'], bp['lpf_stop'])
|
||
wav = _cmb_librosa.resample(_spec_utils.spectrogram_to_wave(spec_s, bp['hl'], mp, d, is_v51_model), orig_sr=bp['sr'], target_sr=sr, res_type=_spec_utils.wav_resolution)
|
||
else:
|
||
if is_v51_model:
|
||
spec_s *= _spec_utils.get_hp_filter_mask(spec_s.shape[1], bp['hpf_start'], bp['hpf_stop'] - 1)
|
||
spec_s *= _spec_utils.get_lp_filter_mask(spec_s.shape[1], bp['lpf_start'], bp['lpf_stop'])
|
||
else:
|
||
spec_s = _spec_utils.fft_hp_filter(spec_s, bp['hpf_start'], bp['hpf_stop'] - 1)
|
||
spec_s = _spec_utils.fft_lp_filter(spec_s, bp['lpf_start'], bp['lpf_stop'])
|
||
wav2 = _cmb_np.add(wav, _spec_utils.spectrogram_to_wave(spec_s, bp['hl'], mp, d, is_v51_model))
|
||
wav = _cmb_librosa.resample(wav2, orig_sr=bp['sr'], target_sr=sr, res_type=_spec_utils.wav_resolution)
|
||
return wav
|
||
_spec_utils.cmb_spectrogram_to_wave = _fixed_cmb_spectrogram_to_wave
|
||
|
||
# 修复 torch.stft/istft 参数兼容性
|
||
import torch as _torch
|
||
_orig_stft = _torch.stft
|
||
def _compat_stft(input, *args, **kwargs):
|
||
names = ['n_fft', 'hop_length', 'win_length', 'window',
|
||
'center', 'pad_mode', 'normalized', 'onesided', 'return_complex']
|
||
for i, v in enumerate(args):
|
||
if i < len(names) and names[i] not in kwargs:
|
||
kwargs[names[i]] = v
|
||
return _orig_stft(input, **kwargs)
|
||
_torch.stft = _compat_stft
|
||
_torch.functional.stft = _compat_stft
|
||
|
||
_uvr_imported = True
|
||
return True
|
||
except Exception as e:
|
||
_uvr_error = f'UVR 模块加载失败: {e}'
|
||
if 'No module named' in str(e):
|
||
_uvr_error += '\n请安装缺少的依赖: pip install pyrubberband ml_collections'
|
||
raise ImportError(_uvr_error) from e
|
||
|
||
|
||
def _secondary_stem_fallback(stem):
|
||
"""当 UVR 未加载时的备用 stem 映射"""
|
||
pairs = {'Vocals': 'Instrumental', 'Instrumental': 'Vocals',
|
||
'Primary Stem': 'Secondary Stem', 'Other': 'No Other',
|
||
'Drums': 'No Drums', 'Bass': 'No Bass', 'Guitar': 'No Guitar'}
|
||
return pairs.get(stem, f'No {stem}' if not stem.startswith('No ') else stem.replace('No ', ''))
|
||
|
||
|
||
def _get_secondary_stem(stem):
|
||
"""获取配对音轨名"""
|
||
if _secondary_stem:
|
||
return _secondary_stem(stem)
|
||
return _secondary_stem_fallback(stem)
|
||
|
||
|
||
# ── 模型元数据 ─────────────────────────────────────────────────────────────────
|
||
|
||
_vr_hash_data = None
|
||
_mdx_hash_data = None
|
||
|
||
|
||
def _load_model_hash_data():
|
||
"""加载 UVR 的模型哈希数据"""
|
||
global _vr_hash_data, _mdx_hash_data
|
||
if _vr_hash_data is not None:
|
||
return
|
||
|
||
cfg = _load_config()
|
||
uvr_path = cfg.get('uvr_project_path', '')
|
||
if not uvr_path:
|
||
_vr_hash_data, _mdx_hash_data = {}, {}
|
||
return
|
||
|
||
vr_dir = os.path.join(uvr_path, 'models', 'VR_Models', 'model_data')
|
||
mdx_dir = os.path.join(uvr_path, 'models', 'MDX_Net_Models', 'model_data')
|
||
|
||
_vr_hash_data = {}
|
||
_mdx_hash_data = {}
|
||
|
||
for path in [
|
||
os.path.join(vr_dir, 'model_data.json'),
|
||
os.path.join(vr_dir, 'model_data_new.json'),
|
||
]:
|
||
if os.path.isfile(path):
|
||
try:
|
||
with open(path, 'r', encoding='utf-8') as f:
|
||
_vr_hash_data.update(json.load(f))
|
||
except Exception:
|
||
pass
|
||
|
||
for path in [
|
||
os.path.join(mdx_dir, 'model_data.json'),
|
||
os.path.join(mdx_dir, 'model_data_new.json'),
|
||
]:
|
||
if os.path.isfile(path):
|
||
try:
|
||
with open(path, 'r', encoding='utf-8') as f:
|
||
_mdx_hash_data.update(json.load(f))
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _compute_model_hash(model_path):
|
||
"""计算模型文件哈希（与 UVR 相同的算法）"""
|
||
try:
|
||
with open(model_path, 'rb') as f:
|
||
f.seek(-10000 * 1024, 2)
|
||
return hashlib.md5(f.read()).hexdigest()
|
||
except Exception:
|
||
try:
|
||
return hashlib.md5(open(model_path, 'rb').read()).hexdigest()
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _lookup_model_meta(model_hash, arch_type):
|
||
"""通过哈希查找模型元数据"""
|
||
_load_model_hash_data()
|
||
if not model_hash:
|
||
return None
|
||
hash_data = _vr_hash_data if 'VR' in arch_type else _mdx_hash_data
|
||
return hash_data.get(model_hash)
|
||
|
||
|
||
# ── HeadlessModelData ─────────────────────────────────────────────────────────
|
||
|
||
class HeadlessModelData:
|
||
"""替代 UVR 的 ModelData，无需 tkinter GUI"""
|
||
|
||
def __init__(self, model_path, process_method, params, model_meta=None):
|
||
self.model_path = model_path
|
||
self.model_name = os.path.splitext(os.path.basename(model_path))[0]
|
||
self.model_basename = self.model_name
|
||
self.process_method = process_method
|
||
self.model_status = True
|
||
self.model_meta = model_meta
|
||
|
||
# GPU 设置
|
||
self.is_gpu_conversion = 0 if params.get('is_gpu', True) else -1
|
||
self.device_set = params.get('device_set', 'Default')
|
||
|
||
# 通用参数
|
||
self.is_normalization = params.get('is_normalization', False)
|
||
self.is_primary_stem_only = params.get('is_primary_stem_only', False)
|
||
self.is_secondary_stem_only = params.get('is_secondary_stem_only', False)
|
||
self.wav_type_set = params.get('wav_type', 'PCM_16')
|
||
self.mp3_bit_set = params.get('mp3_bitrate', '320k')
|
||
self.save_format = params.get('save_format', 'WAV').upper()
|
||
|
||
self.is_invert_spec = False
|
||
self.is_mixer_mode = False
|
||
self.is_mdx_c_seg_def = True
|
||
self.mdx_batch_size = params.get('mdx_batch_size', 1)
|
||
self.mdxnet_stem_select = 'Vocals'
|
||
self.overlap = params.get('mdx_overlap', 0.25)
|
||
self.overlap_mdx = params.get('mdx_overlap', 0.25)
|
||
self.overlap_mdx23 = 8
|
||
self.semitone_shift = 0
|
||
self.is_pitch_change = False
|
||
self.is_match_frequency_pitch = True
|
||
self.is_mdx_combine_stems = False
|
||
self.is_use_opencl = False
|
||
|
||
# MDX 模型数据
|
||
self.is_mdx_ckpt = model_path.endswith('.ckpt') or model_path.endswith('.ckptc')
|
||
self.is_mdx_c = False
|
||
self.mdx_c_configs = None
|
||
self.mdx_model_stems = []
|
||
self.mdx_dim_f_set = None
|
||
self.mdx_dim_t_set = None
|
||
self.mdx_stem_count = 1
|
||
self.compensate = None
|
||
self.mdx_n_fft_scale_set = None
|
||
|
||
# 路径
|
||
paths = _get_uvr_paths()
|
||
self.mixer_path = paths.get('mixer_path', '')
|
||
|
||
# Demucs 默认
|
||
self.demucs_stems = 'All Stems'
|
||
self.is_demucs_combine_stems = False
|
||
self.demucs_source_list = []
|
||
self.demucs_stem_count = 0
|
||
self.demucs_source_map = {}
|
||
self.demucs_version = 'v4'
|
||
|
||
# 主音轨设置
|
||
self.primary_stem = None
|
||
self.secondary_stem = None
|
||
self.primary_stem_native = None
|
||
|
||
# 禁用复杂功能
|
||
self.is_ensemble_mode = False
|
||
self.ensemble_primary_stem = None
|
||
self.ensemble_secondary_stem = None
|
||
self.primary_model_primary_stem = None
|
||
self.is_secondary_model = False
|
||
self.is_secondary_model_activated = False
|
||
self.secondary_model = None
|
||
self.secondary_model_scale = None
|
||
self.pre_proc_model = None
|
||
self.pre_proc_model_activated = False
|
||
self.is_pre_proc_model = False
|
||
self.is_dry_check = False
|
||
self.is_vocal_split_model = False
|
||
self.is_vocal_split_model_activated = False
|
||
self.is_primary_model_primary_stem_only = False
|
||
self.is_primary_model_secondary_stem_only = False
|
||
self.is_save_inst_vocal_splitter = False
|
||
self.is_inst_only_voc_splitter = False
|
||
self.is_save_vocal_only = False
|
||
self.is_deverb_vocals = False
|
||
self.deverb_vocal_opt = 'Vocals'
|
||
self.is_denoise = False
|
||
self.is_denoise_model = False
|
||
self.is_karaoke = False
|
||
self.is_bv_model = False
|
||
self.bv_model_rebalance = 0
|
||
self.is_sec_bv_rebalance = False
|
||
|
||
# 默认
|
||
self.model_samplerate = 44100
|
||
self.model_capacity = (32, 128)
|
||
self.is_vr_51_model = False
|
||
self.is_demucs_pre_proc_model_inst_mix = False
|
||
self.is_change_def = False
|
||
self.is_4_stem_ensemble = False
|
||
self.is_multi_stem_ensemble = False
|
||
self.is_demucs_4_stem_secondaries = False
|
||
self.demucs_4_stem_added_count = 0
|
||
self.model_hash_dir = None
|
||
self.is_get_hash_dir_only = False
|
||
|
||
# 多模型占位
|
||
self.secondary_model_4_stem = []
|
||
self.secondary_model_4_stem_scale = []
|
||
self.secondary_model_4_stem_names = []
|
||
self.secondary_model_4_stem_model_names_list = []
|
||
self.all_models = []
|
||
self.secondary_model_other = None
|
||
self.secondary_model_scale_other = None
|
||
self.secondary_model_bass = None
|
||
self.secondary_model_scale_bass = None
|
||
self.secondary_model_drums = None
|
||
self.secondary_model_scale_drums = None
|
||
|
||
# DeNoise / DeVerb 模型路径
|
||
self.DENOISER_MODEL = paths.get('denoiser_path', '')
|
||
self.DEVERBER_MODEL = paths.get('deverb_path', '')
|
||
self.vocal_split_model = None
|
||
|
||
# 根据架构类型初始化特定参数
|
||
if process_method == VR_ARCH:
|
||
self._init_vr(params, model_meta, paths)
|
||
elif process_method == MDX_ARCH:
|
||
self._init_mdx(params, model_meta, paths)
|
||
elif process_method == DEMUCS_ARCH:
|
||
self._init_demucs(params)
|
||
|
||
def _init_vr(self, params, model_meta, paths):
|
||
"""VR 架构特定参数"""
|
||
self.aggression_setting = float(int(params.get('vr_aggression', 5)) / 100)
|
||
self.is_tta = False
|
||
self.is_post_process = False
|
||
self.window_size = params.get('vr_window_size', 1024)
|
||
self.batch_size = params.get('vr_batch_size', 4)
|
||
self.crop_size = 256
|
||
self.is_high_end_process = 'None'
|
||
self.post_process_threshold = 0.2
|
||
|
||
if model_meta:
|
||
self.primary_stem = model_meta.get('primary_stem', 'Vocals')
|
||
vr_param_name = model_meta.get('vr_model_param', 'bandparam_opposite')
|
||
vr_param_path = os.path.join(paths.get('vr_param_dir', ''), f'{vr_param_name}.json')
|
||
if os.path.isfile(vr_param_path) and _ModelParameters:
|
||
self.vr_model_param = _ModelParameters(vr_param_path)
|
||
self.model_samplerate = self.vr_model_param.param['sr']
|
||
else:
|
||
self.vr_model_param = None
|
||
if 'nout' in model_meta and 'nout_lstm' in model_meta:
|
||
self.model_capacity = (model_meta['nout'], model_meta['nout_lstm'])
|
||
self.is_vr_51_model = True
|
||
else:
|
||
self.vr_model_param = None
|
||
self.primary_stem = 'Vocals'
|
||
|
||
self.primary_stem_native = self.primary_stem
|
||
self.secondary_stem = _get_secondary_stem(self.primary_stem)
|
||
|
||
def _init_mdx(self, params, model_meta, paths):
|
||
"""MDX 架构特定参数"""
|
||
self.margin = 44100
|
||
self.chunks = 0
|
||
seg = params.get('mdx_segment_size', 'Default')
|
||
self.mdx_segment_size = int(seg) if seg and seg != 'Default' else 256
|
||
|
||
if model_meta:
|
||
if 'config_yaml' in model_meta:
|
||
self.is_mdx_c = True
|
||
config_path = os.path.join(
|
||
paths.get('mdx_c_config_path', ''),
|
||
model_meta['config_yaml']
|
||
)
|
||
if os.path.isfile(config_path):
|
||
import yaml
|
||
from ml_collections import ConfigDict
|
||
with open(config_path) as f:
|
||
self.mdx_c_configs = ConfigDict(yaml.load(f, Loader=yaml.FullLoader))
|
||
if self.mdx_c_configs.training.target_instrument:
|
||
target = self.mdx_c_configs.training.target_instrument
|
||
self.mdx_model_stems = [target]
|
||
self.primary_stem = target
|
||
else:
|
||
self.mdx_model_stems = self.mdx_c_configs.training.instruments
|
||
self.mdx_stem_count = len(self.mdx_model_stems)
|
||
self.primary_stem = self.mdx_model_stems[0] if self.mdx_stem_count == 2 else self.mdxnet_stem_select
|
||
else:
|
||
self.primary_stem = model_meta.get('primary_stem', 'Vocals')
|
||
else:
|
||
self.compensate = model_meta.get('compensate', 1.0)
|
||
self.mdx_dim_f_set = model_meta.get('mdx_dim_f_set')
|
||
self.mdx_dim_t_set = model_meta.get('mdx_dim_t_set')
|
||
self.mdx_n_fft_scale_set = model_meta.get('mdx_n_fft_scale_set')
|
||
self.primary_stem = model_meta.get('primary_stem', 'Vocals')
|
||
else:
|
||
self.primary_stem = 'Vocals'
|
||
|
||
self.primary_stem_native = self.primary_stem
|
||
self.secondary_stem = _get_secondary_stem(self.primary_stem)
|
||
|
||
def _init_demucs(self, params):
|
||
"""Demucs 架构特定参数"""
|
||
self.margin_demucs = 44100
|
||
self.chunks_demucs = 0
|
||
self.shifts = 1
|
||
self.is_split_mode = True
|
||
self.segment = params.get('demucs_segment', 'Default')
|
||
self.is_chunk_demucs = False
|
||
self.demucs_stems = params.get('demucs_stems', 'All Stems')
|
||
|
||
# 从文件名推断版本和音轨数
|
||
self.demucs_version = 'v4'
|
||
for ver, tag in [('v1', 'v1 | '), ('v2', 'v2 | '), ('v3', 'v3 | '), ('v4', 'v4 | ')]:
|
||
if tag in self.model_name:
|
||
self.demucs_version = ver
|
||
break
|
||
|
||
# .th 文件名格式为 "sig-checksum"，get_model 只需要 sig 部分
|
||
if self.model_path.endswith('.th') and '-' in self.model_basename:
|
||
self.model_basename = self.model_basename.split('-')[0]
|
||
|
||
if 'UVR_Model' in self.model_name:
|
||
self.demucs_source_list = ['instrumental', 'vocals']
|
||
self.demucs_source_map = {'instrumental': 0, 'vocals': 1}
|
||
self.demucs_stem_count = 2
|
||
self.primary_stem = 'Vocals'
|
||
self.secondary_stem = 'Instrumental'
|
||
else:
|
||
self.demucs_source_list = ['drums', 'bass', 'other', 'vocals']
|
||
self.demucs_source_map = {
|
||
'Bass': 0, 'Drums': 1, 'Other': 2, 'Vocals': 3
|
||
}
|
||
self.demucs_stem_count = 4
|
||
self.primary_stem = 'Vocals'
|
||
self.secondary_stem = _get_secondary_stem('Vocals')
|
||
|
||
|
||
# ── 工具函数 ───────────────────────────────────────────────────────────────────
|
||
|
||
def _update_progress(task_id, step, inference_iterations=0):
|
||
"""更新任务进度 (step: 0.0~1.0)"""
|
||
task = _tasks.get(task_id)
|
||
if not task:
|
||
return
|
||
progress = min(99, max(1, int((step + inference_iterations) * 100)))
|
||
task['progress'] = progress
|
||
|
||
|
||
def _make_process_data(task_id, model_data, audio_path, export_path):
|
||
"""构造 process_data 字典"""
|
||
base = os.path.splitext(os.path.basename(audio_path))[0]
|
||
return {
|
||
'model_data': model_data,
|
||
'export_path': export_path,
|
||
'audio_file_base': base,
|
||
'audio_file': audio_path,
|
||
'set_progress_bar': lambda step, it=0: _update_progress(task_id, step, it),
|
||
'write_to_console': lambda *_, **__: None,
|
||
'process_iteration': lambda: None,
|
||
'cached_source_callback': lambda *_, **__: (None, None),
|
||
'cached_model_source_holder': lambda *_, **__: None,
|
||
'list_all_models': [],
|
||
'is_ensemble_master': False,
|
||
'is_4_stem_ensemble': False,
|
||
}
|
||
|
||
|
||
def _do_separate(task_id, model_path, process_method, audio_path, params, model_meta):
|
||
"""后台线程执行分离"""
|
||
task = _tasks[task_id]
|
||
try:
|
||
_ensure_uvr_imports()
|
||
|
||
import separate as _sep_mod
|
||
import torch
|
||
cuda_ok = getattr(_sep_mod, 'cuda_available', False) and torch.cuda.is_available()
|
||
|
||
model_data = HeadlessModelData(model_path, process_method, params, model_meta)
|
||
|
||
export_path = os.path.join(tempfile.gettempdir(), f'uvr_out_{task_id}')
|
||
os.makedirs(export_path, exist_ok=True)
|
||
|
||
process_data = _make_process_data(task_id, model_data, audio_path, export_path)
|
||
|
||
# 根据架构类型选择分离器
|
||
if process_method == VR_ARCH:
|
||
separator = _SeperateVR(model_data, process_data)
|
||
elif process_method == MDX_ARCH:
|
||
if model_data.is_mdx_c:
|
||
separator = _SeperateMDXC(model_data, process_data)
|
||
else:
|
||
separator = _SeperateMDX(model_data, process_data)
|
||
elif process_method == DEMUCS_ARCH:
|
||
separator = _SeperateDemucs(model_data, process_data)
|
||
else:
|
||
raise ValueError(f'不支持的架构: {process_method}')
|
||
|
||
actual_device = str(getattr(separator, 'device', 'unknown'))
|
||
task['device'] = actual_device
|
||
print(f'[UVR-Sep] is_gpu={params.get("is_gpu")}, cuda={cuda_ok}, device={actual_device}')
|
||
|
||
separator.seperate()
|
||
|
||
# 扫描输出文件
|
||
stems = []
|
||
if os.path.isdir(export_path):
|
||
for fname in sorted(os.listdir(export_path)):
|
||
fpath = os.path.join(export_path, fname)
|
||
if os.path.isfile(fpath):
|
||
stem_name = fname
|
||
if '_(' in fname and fname.endswith(').wav'):
|
||
start = fname.index('_(') + 2
|
||
end = fname.rindex(')')
|
||
stem_name = fname[start:end]
|
||
stems.append({
|
||
'stem': stem_name,
|
||
'filename': fname,
|
||
'path': fpath,
|
||
})
|
||
|
||
if not stems:
|
||
task['status'] = 'error'
|
||
task['error'] = '分离完成但未生成输出文件'
|
||
return
|
||
|
||
task['stems'] = stems
|
||
task['status'] = 'done'
|
||
task['progress'] = 100
|
||
|
||
except Exception as e:
|
||
task['status'] = 'error'
|
||
task['error'] = str(e)
|
||
traceback.print_exc()
|
||
finally:
|
||
try:
|
||
os.remove(audio_path)
|
||
except OSError:
|
||
pass
|
||
|
||
|
||
# ── 路由 ───────────────────────────────────────────────────────────────────────
|
||
|
||
@bp.route('/')
|
||
def page():
|
||
cfg = _load_config()
|
||
return render_template('uvr_sep.html', config=cfg)
|
||
|
||
|
||
@bp.route('/gpu-status')
|
||
def gpu_status():
|
||
"""检测 GPU/CUDA 状态"""
|
||
info = {'cuda_available': False, 'gpu_name': '', 'torch_version': 'N/A', 'torch_cuda': False}
|
||
# 方法1: ctranslate2
|
||
try:
|
||
import ctranslate2
|
||
count = ctranslate2.get_cuda_device_count()
|
||
if count > 0:
|
||
info['cuda_available'] = True
|
||
try:
|
||
info['gpu_name'] = ctranslate2.get_cuda_device_name(0) or ''
|
||
except Exception:
|
||
pass
|
||
except Exception:
|
||
pass
|
||
# 方法2: PyTorch
|
||
try:
|
||
import torch
|
||
info['torch_version'] = torch.__version__
|
||
info['torch_cuda'] = torch.cuda.is_available()
|
||
if info['torch_cuda']:
|
||
info['cuda_available'] = True
|
||
if not info['gpu_name']:
|
||
info['gpu_name'] = torch.cuda.get_device_name(0)
|
||
except ImportError:
|
||
pass
|
||
return jsonify(info)
|
||
|
||
|
||
@bp.route('/config', methods=['GET'])
|
||
def get_config():
|
||
return jsonify(_load_config())
|
||
|
||
|
||
@bp.route('/config', methods=['POST'])
|
||
def save_config():
|
||
data = request.get_json()
|
||
cfg = _load_config()
|
||
for key in DEFAULT_PARAMS:
|
||
if key in data:
|
||
cfg[key] = data[key]
|
||
_save_config(cfg)
|
||
global _uvr_imported, _uvr_error, _vr_hash_data, _mdx_hash_data
|
||
if 'uvr_project_path' in data:
|
||
_uvr_imported = False
|
||
_uvr_error = None
|
||
_vr_hash_data = None
|
||
_mdx_hash_data = None
|
||
return jsonify({'success': True})
|
||
|
||
|
||
@bp.route('/scan-models', methods=['POST'])
|
||
def scan_models():
|
||
"""扫描模型目录"""
|
||
data = request.get_json()
|
||
model_dir = data.get('model_dir', '')
|
||
arch_type = data.get('arch_type', 'MDX-Net')
|
||
|
||
cfg = _load_config()
|
||
model_dir = _get_model_dir_for_arch(arch_type, model_dir, cfg)
|
||
|
||
if not model_dir or not os.path.isdir(model_dir):
|
||
return jsonify({'success': False, 'error': '模型目录无效'}), 400
|
||
|
||
_load_model_hash_data()
|
||
|
||
exts = {
|
||
'VR Arc': ('.pth',),
|
||
'MDX-Net': ('.onnx', '.ckpt', '.ckptc'),
|
||
'Demucs': ('.yaml',),
|
||
}.get(arch_type, ('.onnx', '.ckpt', '.ckptc'))
|
||
|
||
models = []
|
||
try:
|
||
for fname in os.listdir(model_dir):
|
||
fpath = os.path.join(model_dir, fname)
|
||
if not os.path.isfile(fpath):
|
||
continue
|
||
ext = os.path.splitext(fname)[1].lower()
|
||
if ext not in exts:
|
||
continue
|
||
|
||
info = {
|
||
'name': fname,
|
||
'path': fpath,
|
||
'size': os.path.getsize(fpath),
|
||
}
|
||
|
||
model_hash = _compute_model_hash(fpath)
|
||
if model_hash:
|
||
meta = _lookup_model_meta(model_hash, arch_type)
|
||
if meta:
|
||
info['primary_stem'] = meta.get('primary_stem', '')
|
||
info['secondary_stem'] = _get_secondary_stem(meta.get('primary_stem', 'Vocals'))
|
||
info['has_meta'] = True
|
||
else:
|
||
info['has_meta'] = False
|
||
|
||
models.append(info)
|
||
except Exception as e:
|
||
return jsonify({'success': False, 'error': str(e)}), 500
|
||
|
||
return jsonify({'success': True, 'models': models})
|
||
|
||
|
||
@bp.route('/separate', methods=['POST'])
|
||
def start_separate():
|
||
"""开始分离"""
|
||
task_id = request.form.get('task_id', '')
|
||
model_path = request.form.get('model_path', '')
|
||
process_method = request.form.get('process_method', 'Demucs')
|
||
primary_stem = request.form.get('primary_stem', '')
|
||
demucs_stems = request.form.get('demucs_stems', 'All Stems')
|
||
segment = request.form.get('segment', 'Default')
|
||
mdx_overlap = request.form.get('mdx_overlap', '')
|
||
vr_window_size = request.form.get('vr_window_size', '')
|
||
vr_aggression = request.form.get('vr_aggression', '')
|
||
is_primary_stem_only = request.form.get('is_primary_stem_only', '')
|
||
is_secondary_stem_only = request.form.get('is_secondary_stem_only', '')
|
||
is_gpu = request.form.get('is_gpu', '')
|
||
is_normalization = request.form.get('is_normalization', '')
|
||
|
||
if 'audio' not in request.files:
|
||
return jsonify({'success': False, 'error': '请上传音频文件'}), 400
|
||
audio_file = request.files['audio']
|
||
if not audio_file.filename:
|
||
return jsonify({'success': False, 'error': '请上传音频文件'}), 400
|
||
if not model_path:
|
||
return jsonify({'success': False, 'error': '请指定模型路径'}), 400
|
||
if not os.path.isfile(model_path):
|
||
return jsonify({'success': False, 'error': f'模型文件不存在: {model_path}'}), 400
|
||
|
||
ext = os.path.splitext(audio_file.filename)[1].lower()
|
||
if ext not in AUDIO_EXTS:
|
||
return jsonify({'success': False, 'error': f'不支持的音频格式: {ext}'}), 400
|
||
|
||
save_path = os.path.join(tempfile.gettempdir(), f'uvr_{uuid.uuid4().hex}{ext}')
|
||
audio_file.save(save_path)
|
||
|
||
# 架构名直接使用 UVR 常量名
|
||
arch = process_method
|
||
|
||
model_hash = _compute_model_hash(model_path)
|
||
model_meta = _lookup_model_meta(model_hash,
|
||
'VR' if arch == VR_ARCH else 'MDX' if arch == MDX_ARCH else 'Demucs')
|
||
|
||
if not model_meta and primary_stem:
|
||
model_meta = {'primary_stem': primary_stem}
|
||
|
||
if model_meta and primary_stem and primary_stem != model_meta.get('primary_stem', ''):
|
||
model_meta = dict(model_meta)
|
||
model_meta['primary_stem'] = primary_stem
|
||
|
||
if not model_meta and arch != DEMUCS_ARCH:
|
||
try:
|
||
os.remove(save_path)
|
||
except OSError:
|
||
pass
|
||
return jsonify({
|
||
'success': False,
|
||
'error': '无法识别模型，请在模型列表中选择正确的主音轨类型',
|
||
}), 400
|
||
|
||
params = _load_config()
|
||
# 前端传来的主/副音轨选项覆盖配置
|
||
if is_primary_stem_only:
|
||
params['is_primary_stem_only'] = is_primary_stem_only == '1'
|
||
if is_secondary_stem_only:
|
||
params['is_secondary_stem_only'] = is_secondary_stem_only == '1'
|
||
if is_gpu:
|
||
params['is_gpu'] = is_gpu == '1'
|
||
if is_normalization:
|
||
params['is_normalization'] = is_normalization == '1'
|
||
# 架构特异参数覆盖
|
||
if arch == DEMUCS_ARCH:
|
||
params['demucs_stems'] = demucs_stems
|
||
params['demucs_segment'] = segment
|
||
elif arch == MDX_ARCH:
|
||
params['mdx_segment_size'] = segment
|
||
if mdx_overlap:
|
||
try:
|
||
params['mdx_overlap'] = float(mdx_overlap)
|
||
except ValueError:
|
||
pass
|
||
elif arch == VR_ARCH:
|
||
if vr_window_size:
|
||
try:
|
||
params['vr_window_size'] = int(vr_window_size)
|
||
except ValueError:
|
||
pass
|
||
if vr_aggression:
|
||
try:
|
||
params['vr_aggression'] = int(vr_aggression)
|
||
except ValueError:
|
||
pass
|
||
|
||
task_id = task_id or uuid.uuid4().hex
|
||
_tasks[task_id] = {
|
||
'status': 'processing',
|
||
'progress': 0,
|
||
'stems': [],
|
||
'error': None,
|
||
'device': '',
|
||
}
|
||
|
||
threading.Thread(
|
||
target=_do_separate,
|
||
args=(task_id, model_path, arch, save_path, params, model_meta),
|
||
daemon=True,
|
||
).start()
|
||
|
||
return jsonify({'success': True, 'task_id': task_id})
|
||
|
||
|
||
@bp.route('/status/<task_id>')
|
||
def status(task_id):
|
||
task = _tasks.get(task_id)
|
||
if not task:
|
||
return jsonify({'success': False, 'error': '任务不存在'}), 404
|
||
resp = {
|
||
'success': True,
|
||
'status': task['status'],
|
||
'progress': task['progress'],
|
||
}
|
||
if task.get('device'):
|
||
resp['device'] = task['device']
|
||
if task['status'] == 'done':
|
||
resp['stems'] = task['stems']
|
||
resp['count'] = len(task['stems'])
|
||
elif task['status'] == 'error':
|
||
resp['error'] = task.get('error', '未知错误')
|
||
return jsonify(resp)
|
||
|
||
|
||
@bp.route('/download/<task_id>/<stem>')
|
||
def download_stem(task_id, stem):
|
||
task = _tasks.get(task_id)
|
||
if not task or task.get('status') != 'done':
|
||
return jsonify({'error': '文件不存在'}), 404
|
||
for s in task['stems']:
|
||
if s['stem'] == stem and os.path.isfile(s['path']):
|
||
return send_file(s['path'], mimetype='audio/wav',
|
||
as_attachment=True, download_name=s['filename'])
|
||
return jsonify({'error': '文件不存在'}), 404
|
||
|
||
|
||
@bp.route('/download-all/<task_id>')
|
||
def download_all(task_id):
|
||
task = _tasks.get(task_id)
|
||
if not task or task.get('status') != 'done':
|
||
return jsonify({'error': '无文件可下载'}), 404
|
||
|
||
buf = io.BytesIO()
|
||
with zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED) as zf:
|
||
for s in task['stems']:
|
||
if os.path.isfile(s['path']):
|
||
zf.write(s['path'], s['filename'])
|
||
buf.seek(0)
|
||
return send_file(buf, mimetype='application/zip', as_attachment=True,
|
||
download_name='uvr_stems.zip')
|
||
|
||
|
||
@bp.route('/cleanup/<task_id>', methods=['POST'])
|
||
def cleanup(task_id):
|
||
task = _tasks.pop(task_id, None)
|
||
if not task:
|
||
return jsonify({'success': True})
|
||
export_path = os.path.join(tempfile.gettempdir(), f'uvr_out_{task_id}')
|
||
if os.path.isdir(export_path):
|
||
shutil.rmtree(export_path, ignore_errors=True)
|
||
return jsonify({'success': True})
|