Initial commit
This commit is contained in:
王新升
@@ -0,0 +1,146 @@
|
||||
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):
|
||||
self.device = device
|
||||
self.language = language
|
||||
self.save_dir = save_dir
|
||||
self.vocal_sep = vocal_sep
|
||||
self.max_merge_duration = max_merge_duration
|
||||
|
||||
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,
|
||||
)
|
||||
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
|
||||
)
|
||||
|
||||
def run(
|
||||
self,
|
||||
audio_path: str,
|
||||
vocal_sep: bool = True,
|
||||
max_merge_duration: int = 60000,
|
||||
language: str = "Mandarin"
|
||||
) -> 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))
|
||||
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,
|
||||
)
|
||||
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=bool, 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")
|
||||
args = parser.parse_args()
|
||||
|
||||
main(args)
|
||||
Reference in New Issue
Block a user