File size: 7,981 Bytes
e4d8df5 |
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 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 |
import os
import sys
import yaml
import torch
import numpy as np
from pathlib import Path
from hashlib import sha256
sys.path.append(os.getcwd())
from main.configs.config import Config
from main.library.uvr5_separator import spec_utils, common_separator
from main.library.uvr5_separator.demucs import hdemucs, states, apply
translations = Config().translations
sys.path.insert(0, os.path.join(os.getcwd(), "main", "library", "uvr5_separator"))
DEMUCS_4_SOURCE_MAPPER = {common_separator.CommonSeparator.BASS_STEM: 0, common_separator.CommonSeparator.DRUM_STEM: 1, common_separator.CommonSeparator.OTHER_STEM: 2, common_separator.CommonSeparator.VOCAL_STEM: 3}
class DemucsSeparator(common_separator.CommonSeparator):
def __init__(self, common_config, arch_config):
super().__init__(config=common_config)
self.segment_size = arch_config.get("segment_size", "Default")
self.shifts = arch_config.get("shifts", 2)
self.overlap = arch_config.get("overlap", 0.25)
self.segments_enabled = arch_config.get("segments_enabled", True)
self.logger.debug(translations["demucs_info"].format(segment_size=self.segment_size, segments_enabled=self.segments_enabled))
self.logger.debug(translations["demucs_info_2"].format(shifts=self.shifts, overlap=self.overlap))
self.demucs_source_map = DEMUCS_4_SOURCE_MAPPER
self.audio_file_path = None
self.audio_file_base = None
self.demucs_model_instance = None
self.logger.info(translations["start_demucs"])
def separate(self, audio_file_path):
self.logger.debug(translations["start_separator"])
source = None
inst_source = {}
self.audio_file_path = audio_file_path
self.audio_file_base = os.path.splitext(os.path.basename(audio_file_path))[0]
self.logger.debug(translations["prepare_mix"])
mix = self.prepare_mix(self.audio_file_path)
self.logger.debug(translations["demix"].format(shape=mix.shape))
self.logger.debug(translations["cancel_mix"])
self.demucs_model_instance = hdemucs.HDemucs(sources=["drums", "bass", "other", "vocals"])
self.demucs_model_instance = get_demucs_model(name=os.path.splitext(os.path.basename(self.model_path))[0], repo=Path(os.path.dirname(self.model_path)))
self.demucs_model_instance = apply.demucs_segments(self.segment_size, self.demucs_model_instance)
self.demucs_model_instance.to(self.torch_device)
self.demucs_model_instance.eval()
self.logger.debug(translations["model_review"])
source = self.demix_demucs(mix)
del self.demucs_model_instance
self.clear_gpu_cache()
self.logger.debug(translations["del_gpu_cache_after_demix"])
output_files = []
self.logger.debug(translations["process_output_file"])
if isinstance(inst_source, np.ndarray):
self.logger.debug(translations["process_ver"])
inst_source[self.demucs_source_map[common_separator.CommonSeparator.VOCAL_STEM]] = spec_utils.reshape_sources(inst_source[self.demucs_source_map[common_separator.CommonSeparator.VOCAL_STEM]], source[self.demucs_source_map[common_separator.CommonSeparator.VOCAL_STEM]])
source = inst_source
if isinstance(source, np.ndarray):
source_length = len(source)
self.logger.debug(translations["source_length"].format(source_length=source_length))
self.logger.debug(translations["set_map"].format(part=source_length))
match source_length:
case 2: self.demucs_source_map = {common_separator.CommonSeparator.INST_STEM: 0, common_separator.CommonSeparator.VOCAL_STEM: 1}
case 6: self.demucs_source_map = {common_separator.CommonSeparator.BASS_STEM: 0, common_separator.CommonSeparator.DRUM_STEM: 1, common_separator.CommonSeparator.OTHER_STEM: 2, common_separator.CommonSeparator.VOCAL_STEM: 3, common_separator.CommonSeparator.GUITAR_STEM: 4, common_separator.CommonSeparator.PIANO_STEM: 5}
case _: self.demucs_source_map = DEMUCS_4_SOURCE_MAPPER
self.logger.debug(translations["process_all_part"])
for stem_name, stem_value in self.demucs_source_map.items():
if self.output_single_stem is not None:
if stem_name.lower() != self.output_single_stem.lower():
self.logger.debug(translations["skip_part"].format(stem_name=stem_name, output_single_stem=self.output_single_stem))
continue
stem_path = os.path.join(f"{self.audio_file_base}_({stem_name})_{self.model_name}.{self.output_format.lower()}")
self.final_process(stem_path, source[stem_value].T, stem_name)
output_files.append(stem_path)
return output_files
def demix_demucs(self, mix):
self.logger.debug(translations["starting_demix_demucs"])
processed = {}
mix = torch.tensor(mix, dtype=torch.float32)
ref = mix.mean(0)
mix = (mix - ref.mean()) / ref.std()
mix_infer = mix
with torch.no_grad():
self.logger.debug(translations["model_infer"])
sources = apply.apply_model(model=self.demucs_model_instance, mix=mix_infer[None], shifts=self.shifts, split=self.segments_enabled, overlap=self.overlap, static_shifts=1 if self.shifts == 0 else self.shifts, set_progress_bar=None, device=self.torch_device, progress=True)[0]
sources = (sources * ref.std() + ref.mean()).cpu().numpy()
sources[[0, 1]] = sources[[1, 0]]
processed[mix] = sources[:, :, 0:None].copy()
return np.concatenate([s[:, :, 0:None] for s in list(processed.values())], axis=-1)
class LocalRepo:
def __init__(self, root):
self.root = root
self.scan()
def scan(self):
self._models, self._checksums = {}, {}
for file in self.root.iterdir():
if file.suffix == ".th":
if "-" in file.stem:
xp_sig, checksum = file.stem.split("-")
self._checksums[xp_sig] = checksum
else: xp_sig = file.stem
if xp_sig in self._models: raise RuntimeError(translations["del_all_but_one"].format(xp_sig=xp_sig))
self._models[xp_sig] = file
def has_model(self, sig):
return sig in self._models
def get_model(self, sig):
try:
file = self._models[sig]
except KeyError:
raise RuntimeError(translations["not_found_model_signature"].format(sig=sig))
if sig in self._checksums: check_checksum(file, self._checksums[sig])
return states.load_model(file)
class BagOnlyRepo:
def __init__(self, root, model_repo):
self.root = root
self.model_repo = model_repo
self.scan()
def scan(self):
self._bags = {}
for file in self.root.iterdir():
if file.suffix == ".yaml": self._bags[file.stem] = file
def get_model(self, name):
try:
yaml_file = self._bags[name]
except KeyError:
raise RuntimeError(translations["name_not_pretrained"].format(name=name))
bag = yaml.safe_load(open(yaml_file))
return apply.BagOfModels([self.model_repo.get_model(sig) for sig in bag["models"]], bag.get("weights"), bag.get("segment"))
def check_checksum(path, checksum):
sha = sha256()
with open(path, "rb") as file:
while 1:
buf = file.read(2**20)
if not buf: break
sha.update(buf)
actual_checksum = sha.hexdigest()[: len(checksum)]
if actual_checksum != checksum: raise RuntimeError(translations["invalid_checksum"].format(path=path, checksum=checksum, actual_checksum=actual_checksum))
def get_demucs_model(name, repo = None):
model_repo = LocalRepo(repo)
return (model_repo.get_model(name) if model_repo.has_model(name) else BagOnlyRepo(repo, model_repo).get_model(name)).eval() |