162 lines
6.6 KiB
Python
162 lines
6.6 KiB
Python
import json
|
|
import shutil
|
|
import soundfile as sf
|
|
from pathlib import Path
|
|
import librosa
|
|
|
|
from preprocess.utils import convert_metadata, merge_short_segments
|
|
|
|
from preprocess.tools import (
|
|
F0Extractor,
|
|
VocalDetector,
|
|
VocalSeparator,
|
|
NoteTranscriber,
|
|
LyricTranscriber,
|
|
)
|
|
|
|
|
|
class PreprocessPipeline:
|
|
def __init__(self, device: str, language: str, save_dir: str, vocal_sep: bool = True, max_merge_duration: int = 60000, midi_transcribe: bool = True):
|
|
self.device = device
|
|
self.language = language
|
|
self.save_dir = save_dir
|
|
self.vocal_sep = vocal_sep
|
|
self.max_merge_duration = max_merge_duration
|
|
self.midi_transcribe = midi_transcribe
|
|
|
|
if vocal_sep:
|
|
self.vocal_separator = VocalSeparator(
|
|
sep_model_path="pretrained_models/SoulX-Singer-Preprocess/mel-band-roformer-karaoke/mel_band_roformer_karaoke_becruily.ckpt",
|
|
sep_config_path="pretrained_models/SoulX-Singer-Preprocess/mel-band-roformer-karaoke/config_karaoke_becruily.yaml",
|
|
der_model_path="pretrained_models/SoulX-Singer-Preprocess/dereverb_mel_band_roformer/dereverb_mel_band_roformer_anvuew_sdr_19.1729.ckpt",
|
|
der_config_path="pretrained_models/SoulX-Singer-Preprocess/dereverb_mel_band_roformer/dereverb_mel_band_roformer_anvuew.yaml",
|
|
device=device
|
|
)
|
|
else:
|
|
self.vocal_separator = None
|
|
self.f0_extractor = F0Extractor(
|
|
model_path="pretrained_models/SoulX-Singer-Preprocess/rmvpe/rmvpe.pt",
|
|
device=device,
|
|
)
|
|
if self.midi_transcribe:
|
|
self.vocal_detector = VocalDetector(
|
|
cut_wavs_output_dir= f"{save_dir}/cut_wavs",
|
|
)
|
|
self.lyric_transcriber = LyricTranscriber(
|
|
zh_model_path="pretrained_models/SoulX-Singer-Preprocess/speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch",
|
|
en_model_path="pretrained_models/SoulX-Singer-Preprocess/parakeet-tdt-0.6b-v2/parakeet-tdt-0.6b-v2.nemo",
|
|
device=device
|
|
)
|
|
self.note_transcriber = NoteTranscriber(
|
|
rosvot_model_path="pretrained_models/SoulX-Singer-Preprocess/rosvot/rosvot/model.pt",
|
|
rwbd_model_path="pretrained_models/SoulX-Singer-Preprocess/rosvot/rwbd/model.pt",
|
|
device=device
|
|
)
|
|
else:
|
|
self.vocal_detector = None
|
|
self.lyric_transcriber = None
|
|
self.note_transcriber = None
|
|
|
|
def run(
|
|
self,
|
|
audio_path: str,
|
|
vocal_sep: bool = None,
|
|
max_merge_duration: int = None,
|
|
language: str = None,
|
|
) -> None:
|
|
vocal_sep = self.vocal_sep if vocal_sep is None else vocal_sep
|
|
max_merge_duration = self.max_merge_duration if max_merge_duration is None else max_merge_duration
|
|
language = self.language if language is None else language
|
|
output_dir = Path(self.save_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
if vocal_sep:
|
|
# Perform vocal/accompaniment separation
|
|
sep = self.vocal_separator.process(audio_path)
|
|
vocal = sep.vocals_dereverbed.T
|
|
acc = sep.accompaniment.T
|
|
sample_rate = sep.sample_rate
|
|
|
|
vocal_path = output_dir / "vocal.wav"
|
|
acc_path = output_dir / "acc.wav"
|
|
sf.write(vocal_path, vocal, sample_rate)
|
|
sf.write(acc_path, acc, sample_rate)
|
|
else:
|
|
# Use the original audio as vocal source (no separation)
|
|
vocal, sample_rate = librosa.load(audio_path, sr=None, mono=True)
|
|
vocal_path = output_dir / "vocal.wav"
|
|
sf.write(vocal_path, vocal, sample_rate)
|
|
|
|
vocal_f0 = self.f0_extractor.process(str(vocal_path), f0_path=str(vocal_path).replace(".wav", "_f0.npy"))
|
|
|
|
if not self.midi_transcribe or self.vocal_detector is None or self.lyric_transcriber is None or self.note_transcriber is None:
|
|
return
|
|
|
|
segments = self.vocal_detector.process(str(vocal_path), f0=vocal_f0)
|
|
|
|
metadata = []
|
|
for seg in segments:
|
|
self.f0_extractor.process(seg["wav_fn"], f0_path=seg["wav_fn"].replace(".wav", "_f0.npy"))
|
|
words, durs = self.lyric_transcriber.process(
|
|
seg["wav_fn"], language
|
|
)
|
|
seg["words"] = words
|
|
seg["word_durs"] = durs
|
|
seg["language"] = language
|
|
metadata.append(
|
|
self.note_transcriber.process(seg, segment_info=seg)
|
|
)
|
|
|
|
merged = merge_short_segments(
|
|
vocal,
|
|
sample_rate,
|
|
metadata,
|
|
output_dir / "long_cut_wavs",
|
|
max_duration_ms=max_merge_duration,
|
|
)
|
|
|
|
final_metadata = []
|
|
|
|
for item in merged:
|
|
self.f0_extractor.process(item.wav_fn, f0_path=item.wav_fn.replace(".wav", "_f0.npy"))
|
|
final_metadata.append(convert_metadata(item))
|
|
|
|
with open(output_dir / "metadata.json", "w", encoding="utf-8") as f:
|
|
json.dump(final_metadata, f, ensure_ascii=False, indent=2)
|
|
|
|
shutil.copy(output_dir / "metadata.json", audio_path.replace(".wav", ".json").replace(".mp3", ".json").replace(".flac", ".json"))
|
|
|
|
|
|
def main(args):
|
|
pipeline = PreprocessPipeline(
|
|
device=args.device,
|
|
language=args.language,
|
|
save_dir=args.save_dir,
|
|
vocal_sep=args.vocal_sep,
|
|
max_merge_duration=args.max_merge_duration,
|
|
midi_transcribe=args.midi_transcribe,
|
|
)
|
|
pipeline.run(
|
|
audio_path=args.audio_path,
|
|
language=args.language,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--audio_path", type=str, required=True, help="Path to the input audio file")
|
|
parser.add_argument("--save_dir", type=str, required=True, help="Directory to save the output files")
|
|
parser.add_argument("--language", type=str, default="Mandarin", help="Language of the audio")
|
|
parser.add_argument("--device", type=str, default="cuda:0", help="Device to run the models on")
|
|
parser.add_argument("--vocal_sep", type=str, default="True", help="Whether to perform vocal separation")
|
|
parser.add_argument("--max_merge_duration", type=int, default=60000, help="Maximum merged segment duration in milliseconds")
|
|
parser.add_argument("--midi_transcribe", type=str, default="True", help="Whether to do MIDI transcription")
|
|
args = parser.parse_args()
|
|
|
|
args.vocal_sep = args.vocal_sep.lower() == "true"
|
|
args.midi_transcribe = args.midi_transcribe.lower() == "true"
|
|
|
|
main(args)
|